diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 340a3ba1a..8c3afae4a 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -11,7 +11,7 @@ const COPILOT_PREMIUM_MULTIPLIERS: Record = { import * as path from "node:path"; import { $env } from "@oh-my-pi/pi-utils"; -import { AuthCredentialStore } from "../src/auth-storage"; +import { SqliteAuthCredentialStore } from "../src/auth-storage"; import { createModelManager } from "../src/model-manager"; import { applyGeneratedModelPolicies, @@ -51,7 +51,7 @@ async function resolveProviderApiKey(providerId: string, catalog: CatalogDiscove } try { - const storage = await AuthCredentialStore.open(); + const storage = await SqliteAuthCredentialStore.open(); try { const storedApiKey = storage.getApiKey(providerId); if (storedApiKey) { @@ -214,7 +214,7 @@ const ANTIGRAVITY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com" async function getOAuthCredentialsFromStorage(provider: OAuthProvider): Promise { try { - const storage = await AuthCredentialStore.open(); + const storage = await SqliteAuthCredentialStore.open(); try { const creds = storage.getOAuth(provider); if (!creds) { diff --git a/packages/ai/src/auth-broker/client.ts b/packages/ai/src/auth-broker/client.ts new file mode 100644 index 000000000..5ea6eba60 --- /dev/null +++ b/packages/ai/src/auth-broker/client.ts @@ -0,0 +1,135 @@ +/** + * HTTP client for the omp auth-broker server. + * + * Used by {@link RemoteAuthCredentialStore} (snapshot pulls) and by + * `omp auth-broker status` (liveness checks). All endpoints except + * `/v1/healthz` require a bearer token. + */ +import type { AuthCredential } from "../auth-storage"; +import type { + CredentialDisableRequest, + CredentialDisableResponse, + CredentialRefreshResponse, + CredentialUploadRequest, + CredentialUploadResponse, + HealthzResponse, + SnapshotResponse, +} from "./types"; + +export interface AuthBrokerClientOptions { + /** Base URL (e.g. `https://broker.tailnet:8765`). Trailing slashes are trimmed. */ + url: string; + /** Bearer token used for everything except `healthz`. */ + token: string; + /** Per-request timeout in milliseconds. Default 10s. */ + timeoutMs?: number; + /** Retry connection errors this many times. Default 1. */ + maxRetries?: number; + /** Override fetch (used in tests). Default global `fetch`. */ + fetchImpl?: typeof fetch; +} + +export class AuthBrokerError extends Error { + readonly status: number | undefined; + readonly body: string | undefined; + constructor(message: string, opts: { status?: number; body?: string; cause?: unknown } = {}) { + super(message, { cause: opts.cause }); + this.name = "AuthBrokerError"; + this.status = opts.status; + this.body = opts.body; + } +} + +const DEFAULT_TIMEOUT_MS = 10_000; +const DEFAULT_MAX_RETRIES = 1; + +export class AuthBrokerClient { + readonly #baseUrl: string; + readonly #token: string; + readonly #timeoutMs: number; + readonly #maxRetries: number; + readonly #fetch: typeof fetch; + + constructor(opts: AuthBrokerClientOptions) { + this.#baseUrl = opts.url.replace(/\/+$/, ""); + this.#token = opts.token; + this.#timeoutMs = opts.timeoutMs ?? DEFAULT_TIMEOUT_MS; + this.#maxRetries = opts.maxRetries ?? DEFAULT_MAX_RETRIES; + this.#fetch = opts.fetchImpl ?? fetch; + } + + healthz(): Promise { + return this.#request("GET", "/v1/healthz", { auth: false }); + } + + fetchSnapshot(): Promise { + return this.#request("GET", "/v1/snapshot"); + } + + async refreshCredential(id: number): Promise { + return this.#request("POST", `/v1/credential/${id}/refresh`); + } + + async disableCredential(id: number, cause: string): Promise { + const body: CredentialDisableRequest = { cause }; + return this.#request("POST", `/v1/credential/${id}/disable`, { + body, + }); + } + + async uploadCredential(provider: string, credential: AuthCredential): Promise { + const body: CredentialUploadRequest = { provider, credential }; + return this.#request("POST", "/v1/credential", { body }); + } + + async #request(method: "GET" | "POST", path: string, opts: { auth?: boolean; body?: unknown } = {}): Promise { + const auth = opts.auth ?? true; + const url = `${this.#baseUrl}${path}`; + const headers: Record = { Accept: "application/json" }; + if (auth) headers.Authorization = `Bearer ${this.#token}`; + let payload: string | undefined; + if (opts.body !== undefined) { + payload = JSON.stringify(opts.body); + headers["Content-Type"] = "application/json"; + } + + let lastError: unknown; + for (let attempt = 0; attempt <= this.#maxRetries; attempt += 1) { + try { + const response = await this.#fetch(url, { + method, + headers, + body: payload, + signal: AbortSignal.timeout(this.#timeoutMs), + }); + const text = await response.text(); + if (!response.ok) { + throw new AuthBrokerError(`Auth broker request failed: ${response.status} ${response.statusText}`, { + status: response.status, + body: text, + }); + } + if (!text) return undefined as T; + try { + return JSON.parse(text) as T; + } catch (parseError) { + throw new AuthBrokerError("Auth broker returned malformed JSON", { + status: response.status, + body: text, + cause: parseError, + }); + } + } catch (error) { + lastError = error; + if (error instanceof AuthBrokerError && error.status !== undefined) { + // HTTP errors (4xx/5xx) don't retry — caller knows what to do. + throw error; + } + if (attempt >= this.#maxRetries) break; + } + } + throw new AuthBrokerError(`Auth broker request failed after ${this.#maxRetries + 1} attempt(s)`, { + cause: lastError, + }); + } +} diff --git a/packages/ai/src/auth-broker/index.ts b/packages/ai/src/auth-broker/index.ts new file mode 100644 index 000000000..4858fbfdf --- /dev/null +++ b/packages/ai/src/auth-broker/index.ts @@ -0,0 +1,5 @@ +export * from "./client"; +export * from "./refresher"; +export * from "./remote-store"; +export * from "./server"; +export * from "./types"; diff --git a/packages/ai/src/auth-broker/refresher.ts b/packages/ai/src/auth-broker/refresher.ts new file mode 100644 index 000000000..7b9e9947b --- /dev/null +++ b/packages/ai/src/auth-broker/refresher.ts @@ -0,0 +1,116 @@ +/** + * Background OAuth refresh loop for the auth-broker server. + * + * Iterates active OAuth credentials at `refreshIntervalMs` cadence, refreshing + * any whose `expires - Date.now() < refreshSkewMs`. Single-flighted per + * credential id so a long refresh can't be retriggered until it settles. + * + * Definitively-failed credentials (invalid_grant / 401 not from network blip) + * are disabled via {@link AuthStorage.disableCredentialById} so the next + * snapshot pull surfaces a clean delete on the client. + */ +import { logger } from "@oh-my-pi/pi-utils"; +import type { AuthStorage } from "../auth-storage"; +import { DEFAULT_REFRESH_INTERVAL_MS, DEFAULT_REFRESH_SKEW_MS } from "./types"; + +export interface AuthBrokerRefresherOptions { + storage: AuthStorage; + /** Refresh credentials expiring within this window. Default 5 min. */ + refreshSkewMs?: number; + /** Loop cadence. Default 60s. */ + refreshIntervalMs?: number; + /** Override clock (tests). */ + now?: () => number; +} + +const INVALID_GRANT_REGEX = /invalid_grant|invalid_token|revoked|unauthorized|expired.*refresh|refresh.*expired/i; +const TRANSIENT_REGEX = /timeout|network|fetch failed|ECONNREFUSED/i; +const HTTP_401_403_REGEX = /\b(401|403)\b/; + +function isDefinitiveFailure(errorMsg: string): boolean { + if (INVALID_GRANT_REGEX.test(errorMsg)) return true; + if (HTTP_401_403_REGEX.test(errorMsg) && !TRANSIENT_REGEX.test(errorMsg)) return true; + return false; +} + +export class AuthBrokerRefresher { + readonly #storage: AuthStorage; + readonly #refreshSkewMs: number; + readonly #refreshIntervalMs: number; + readonly #now: () => number; + readonly #inFlight: Map> = new Map(); + #timer: NodeJS.Timeout | undefined; + #running = false; + + constructor(opts: AuthBrokerRefresherOptions) { + this.#storage = opts.storage; + this.#refreshSkewMs = opts.refreshSkewMs ?? DEFAULT_REFRESH_SKEW_MS; + this.#refreshIntervalMs = opts.refreshIntervalMs ?? DEFAULT_REFRESH_INTERVAL_MS; + this.#now = opts.now ?? Date.now; + } + + start(): void { + if (this.#timer !== undefined) return; + // Refresh sweep is best-effort; kick once immediately so freshly-booted + // brokers don't hand out near-expired tokens for the first interval. + void this.tick(); + this.#timer = setInterval(() => { + void this.tick(); + }, this.#refreshIntervalMs); + } + + stop(): void { + if (this.#timer !== undefined) { + clearInterval(this.#timer); + this.#timer = undefined; + } + } + + /** Run one sweep. Exposed for tests. */ + async tick(): Promise { + if (this.#running) return; + this.#running = true; + try { + await this.#storage.reload(); + const snapshot = this.#storage.exportSnapshot(); + const now = this.#now(); + const deadline = now + this.#refreshSkewMs; + const targets: number[] = []; + for (const entry of snapshot.credentials) { + if (entry.credential.type !== "oauth") continue; + const expires = entry.credential.expires; + if (typeof expires !== "number" || !Number.isFinite(expires)) continue; + if (expires > deadline) continue; + targets.push(entry.id); + } + await Promise.all(targets.map(id => this.#refreshOne(id))); + } finally { + this.#running = false; + } + } + + #refreshOne(id: number): Promise { + const existing = this.#inFlight.get(id); + if (existing) return existing; + const promise = (async () => { + try { + await this.#storage.forceRefreshCredentialById(id); + } catch (error) { + const errorMsg = String(error); + if (isDefinitiveFailure(errorMsg)) { + logger.warn("auth-broker refresh failed definitively; disabling credential", { + id, + error: errorMsg, + }); + this.#storage.disableCredentialById(id, `auth-broker refresh failed: ${errorMsg}`); + } else { + logger.debug("auth-broker refresh failed (transient)", { id, error: errorMsg }); + } + } finally { + this.#inFlight.delete(id); + } + })(); + this.#inFlight.set(id, promise); + return promise; + } +} diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts new file mode 100644 index 000000000..f4b3b46e6 --- /dev/null +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -0,0 +1,144 @@ +/** + * Client-side {@link AuthCredentialStore} that mirrors a remote broker's + * snapshot. Refresh tokens never leave the broker; mutating methods (`replace*`, + * `upsert*`, `delete*ForProvider`) throw because login flows are server-side. + * + * Cache (`getCache`/`setCache`/`cleanExpiredCache`) is in-memory and ephemeral — + * usage reports cache TTL is ~30s, so durability across runs isn't required. + */ +import { logger } from "@oh-my-pi/pi-utils"; +import type { + AuthCredential, + AuthCredentialSnapshot, + AuthCredentialStore, + StoredAuthCredential, +} from "../auth-storage"; +import type { AuthBrokerClient } from "./client"; + +interface CacheEntry { + value: string; + expiresAtSec: number; +} + +export interface RemoteAuthCredentialStoreOptions { + client: AuthBrokerClient; + /** + * Initial snapshot. When omitted, callers must call + * {@link RemoteAuthCredentialStore.refreshSnapshot} before the first read. + */ + initialSnapshot?: AuthCredentialSnapshot; +} + +export class RemoteAuthCredentialStore implements AuthCredentialStore { + readonly #client: AuthBrokerClient; + #snapshot: AuthCredentialSnapshot; + #cache: Map = new Map(); + #closed = false; + + constructor(opts: RemoteAuthCredentialStoreOptions) { + this.#client = opts.client; + this.#snapshot = opts.initialSnapshot ?? { generatedAt: 0, credentials: [] }; + } + + get client(): AuthBrokerClient { + return this.#client; + } + + get snapshot(): AuthCredentialSnapshot { + return this.#snapshot; + } + + /** Re-hydrate the in-memory snapshot from the broker. */ + async refreshSnapshot(): Promise { + this.#snapshot = await this.#client.fetchSnapshot(); + return this.#snapshot; + } + + listAuthCredentials(provider?: string): StoredAuthCredential[] { + const out: StoredAuthCredential[] = []; + for (const entry of this.#snapshot.credentials) { + if (provider !== undefined && entry.provider !== provider) continue; + out.push({ + id: entry.id, + provider: entry.provider, + credential: entry.credential as AuthCredential, + disabledCause: null, + }); + } + return out; + } + + /** + * In-memory update from a successful refresh through the broker. AuthStorage + * calls this after `#replaceCredentialAt`; the broker already persisted the + * authoritative row, so we just mirror it. + */ + updateAuthCredential(id: number, credential: AuthCredential): void { + for (const entry of this.#snapshot.credentials) { + if (entry.id !== id) continue; + entry.credential = credential as typeof entry.credential; + return; + } + } + + deleteAuthCredential(id: number, disabledCause: string): void { + const next = this.#snapshot.credentials.filter(entry => entry.id !== id); + this.#snapshot = { ...this.#snapshot, credentials: next }; + // Fire-and-forget: tell the broker to persist the disable. + this.#client.disableCredential(id, disabledCause).catch(error => { + logger.warn("auth-broker disable propagation failed", { id, error: String(error) }); + }); + } + + tryDisableAuthCredentialIfMatches(id: number, _expectedData: string, disabledCause: string): boolean { + const found = this.#snapshot.credentials.find(entry => entry.id === id); + if (!found) return false; + this.deleteAuthCredential(id, disabledCause); + return true; + } + + replaceAuthCredentialsForProvider(_provider: string, _credentials: AuthCredential[]): StoredAuthCredential[] { + throw new Error( + "RemoteAuthCredentialStore is read-only on the client. Use `omp auth-broker login ` to mutate credentials.", + ); + } + + upsertAuthCredentialForProvider(_provider: string, _credential: AuthCredential): StoredAuthCredential[] { + throw new Error( + "RemoteAuthCredentialStore is read-only on the client. Use `omp auth-broker login ` to mutate credentials.", + ); + } + + deleteAuthCredentialsForProvider(_provider: string, _disabledCause: string): void { + throw new Error( + "RemoteAuthCredentialStore is read-only on the client. Use `omp auth-broker logout ` to mutate credentials.", + ); + } + + getCache(key: string): string | null { + const entry = this.#cache.get(key); + if (!entry) return null; + if (entry.expiresAtSec * 1000 <= Date.now()) { + this.#cache.delete(key); + return null; + } + return entry.value; + } + + setCache(key: string, value: string, expiresAtSec: number): void { + this.#cache.set(key, { value, expiresAtSec }); + } + + cleanExpiredCache(): void { + const nowSec = Math.floor(Date.now() / 1000); + for (const [key, entry] of this.#cache) { + if (entry.expiresAtSec <= nowSec) this.#cache.delete(key); + } + } + + close(): void { + if (this.#closed) return; + this.#closed = true; + this.#cache.clear(); + } +} diff --git a/packages/ai/src/auth-broker/server.ts b/packages/ai/src/auth-broker/server.ts new file mode 100644 index 000000000..cd2054438 --- /dev/null +++ b/packages/ai/src/auth-broker/server.ts @@ -0,0 +1,233 @@ +/** + * Auth broker HTTP server. + * + * Wraps an {@link AuthStorage} (backed by a SQLite store on the broker host) + * and exposes a minimal REST API for snapshot pulls and explicit refresh / + * disable operations. Background refresh of expiring credentials lives in + * {@link AuthBrokerRefresher}. + * + * Transport security is delegated to the operator (Tailscale / Wireguard); + * the server only checks a bearer token against an allow-list per request. + */ +import { logger } from "@oh-my-pi/pi-utils"; +import type { AuthStorage } from "../auth-storage"; +import { AuthBrokerRefresher } from "./refresher"; +import type { + CredentialDisableRequest, + CredentialDisableResponse, + CredentialRefreshResponse, + CredentialUploadRequest, + CredentialUploadResponse, + HealthzResponse, + SnapshotResponse, +} from "./types"; +import { DEFAULT_AUTH_BROKER_BIND, DEFAULT_REFRESH_INTERVAL_MS, DEFAULT_REFRESH_SKEW_MS } from "./types"; + +export interface AuthBrokerServerOptions { + /** Underlying credential storage (wraps the local SQLite store on the broker). */ + storage: AuthStorage; + /** Listen address; accepts `host:port` or just `port`. */ + bind?: string; + /** Accept any of these bearer tokens. Empty disables auth (loopback only). */ + bearerTokens: string[]; + /** Broker version string surfaced on `/v1/healthz`. */ + version?: string; + /** Refresh credentials expiring within this window. Default 5 min. */ + refreshSkewMs?: number; + /** Background refresh cadence. Default 60s. */ + refreshIntervalMs?: number; + /** Disable the background refresher (e.g. for tests). */ + disableRefresher?: boolean; +} + +export interface AuthBrokerServerHandle { + /** Bound URL (`http://host:port`). */ + url: string; + port: number; + hostname: string; + close(): Promise; +} + +interface ParsedBind { + hostname: string; + port: number; +} + +function parseBind(raw: string): ParsedBind { + const trimmed = raw.trim(); + if (/^\d+$/.test(trimmed)) { + return { hostname: "127.0.0.1", port: Number.parseInt(trimmed, 10) }; + } + const lastColon = trimmed.lastIndexOf(":"); + if (lastColon < 0) { + throw new Error(`Invalid bind '${raw}'; expected 'host:port' or 'port'.`); + } + const hostPart = trimmed.slice(0, lastColon); + const portPart = trimmed.slice(lastColon + 1); + const port = Number.parseInt(portPart, 10); + if (!Number.isFinite(port) || port < 0 || port > 65535) { + throw new Error(`Invalid bind '${raw}'; port out of range.`); + } + return { hostname: hostPart, port }; +} + +function json(status: number, body: unknown): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +function isAuthorized(req: Request, tokens: ReadonlySet): boolean { + if (tokens.size === 0) return true; + const header = req.headers.get("authorization"); + if (!header) return false; + const match = header.match(/^Bearer\s+(.+)$/i); + if (!match) return false; + return tokens.has(match[1].trim()); +} + +const REFRESH_ROUTE = /^\/v1\/credential\/(\d+)\/refresh$/; +const DISABLE_ROUTE = /^\/v1\/credential\/(\d+)\/disable$/; + +/** Boot the broker. Caller owns lifecycle; `handle.close()` to stop. */ +export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServerHandle { + const bind = parseBind(opts.bind ?? DEFAULT_AUTH_BROKER_BIND); + const tokens = new Set(opts.bearerTokens); + const version = opts.version; + + const refresher = opts.disableRefresher + ? undefined + : new AuthBrokerRefresher({ + storage: opts.storage, + refreshSkewMs: opts.refreshSkewMs ?? DEFAULT_REFRESH_SKEW_MS, + refreshIntervalMs: opts.refreshIntervalMs ?? DEFAULT_REFRESH_INTERVAL_MS, + }); + refresher?.start(); + + const server = Bun.serve({ + hostname: bind.hostname, + port: bind.port, + fetch: async (req): Promise => { + const url = new URL(req.url); + const pathname = url.pathname; + const peer = + req.headers.get("x-forwarded-for")?.split(",")[0].trim() || req.headers.get("x-real-ip") || "unknown"; + try { + if (req.method === "GET" && pathname === "/v1/healthz") { + const body: HealthzResponse = { ok: true, version }; + return json(200, body); + } + if (!isAuthorized(req, tokens)) { + logger.info("auth-broker request unauthorized", { method: req.method, path: pathname, peer }); + return json(401, { error: "unauthorized" }); + } + if (req.method === "GET" && pathname === "/v1/snapshot") { + await opts.storage.reload(); + const body: SnapshotResponse = opts.storage.exportSnapshot(); + logger.info("auth-broker snapshot served", { peer, credentials: body.credentials.length }); + return json(200, body); + } + const refreshMatch = req.method === "POST" ? pathname.match(REFRESH_ROUTE) : null; + if (refreshMatch) { + const id = Number.parseInt(refreshMatch[1], 10); + try { + const entry = await opts.storage.forceRefreshCredentialById(id); + const body: CredentialRefreshResponse = { entry }; + logger.info("auth-broker credential refreshed", { + id, + provider: entry.provider, + peer, + expires: entry.credential.type === "oauth" ? entry.credential.expires : undefined, + }); + return json(200, body); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn("auth-broker refresh failed", { id, peer, error: message }); + const status = message.includes("No credential with id") ? 404 : 500; + return json(status, { error: message }); + } + } + const disableMatch = req.method === "POST" ? pathname.match(DISABLE_ROUTE) : null; + if (disableMatch) { + const id = Number.parseInt(disableMatch[1], 10); + let cause = "disabled via auth-broker"; + try { + const body = (await req.json()) as Partial; + if (typeof body?.cause === "string" && body.cause.length > 0) cause = body.cause; + } catch { + // Empty / malformed body — default cause already set. + } + const ok = opts.storage.disableCredentialById(id, cause); + if (!ok) { + logger.info("auth-broker disable miss", { id, peer, cause }); + return json(404, { error: `No credential with id=${id}` }); + } + logger.info("auth-broker credential disabled", { id, peer, cause }); + const response: CredentialDisableResponse = { ok: true }; + return json(200, response); + } + if (req.method === "POST" && pathname === "/v1/credential") { + let body: Partial; + try { + body = (await req.json()) as Partial; + } catch (error) { + return json(400, { error: `Invalid JSON body: ${String(error)}` }); + } + if (!body || typeof body.provider !== "string" || body.provider.length === 0) { + return json(400, { error: "Missing `provider` field" }); + } + if (!body.credential || typeof body.credential !== "object") { + return json(400, { error: "Missing `credential` field" }); + } + const credential = body.credential; + if (credential.type !== "oauth" && credential.type !== "api_key") { + return json(400, { + error: `Invalid credential.type: ${String((credential as { type?: unknown }).type)}`, + }); + } + try { + const entries = opts.storage.upsertCredential(body.provider, credential); + const identity = + credential.type === "oauth" + ? (credential.email ?? credential.accountId ?? credential.projectId ?? "(no identity)") + : "(api key)"; + logger.info("auth-broker credential upserted", { + provider: body.provider, + type: credential.type, + identity, + peer, + providerTotal: entries.length, + }); + const response: CredentialUploadResponse = { entries }; + return json(200, response); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn("auth-broker upload failed", { provider: body.provider, peer, error: message }); + return json(500, { error: message }); + } + } + return json(404, { error: `No route: ${req.method} ${pathname}` }); + } catch (error) { + logger.error("auth-broker handler crashed", { + method: req.method, + path: pathname, + error: String(error), + }); + return json(500, { error: "internal error" }); + } + }, + }); + + const boundHost = server.hostname ?? bind.hostname; + const boundPort = server.port ?? bind.port; + return { + url: `http://${boundHost}:${boundPort}`, + port: boundPort, + hostname: boundHost, + close: async () => { + refresher?.stop(); + server.stop(true); + }, + }; +} diff --git a/packages/ai/src/auth-broker/types.ts b/packages/ai/src/auth-broker/types.ts new file mode 100644 index 000000000..59a04efea --- /dev/null +++ b/packages/ai/src/auth-broker/types.ts @@ -0,0 +1,62 @@ +/** + * Wire types shared between the auth-broker server and clients. + * + * The broker holds OAuth refresh tokens and exposes a redacted snapshot; + * clients use `access` tokens directly and call back to the broker when a + * credential expires or a 401 surfaces on a supposedly-fresh credential. + */ + +import type { AuthCredential, AuthCredentialSnapshot, AuthCredentialSnapshotEntry } from "../auth-storage"; + +/** GET /v1/healthz response body. */ +export interface HealthzResponse { + ok: boolean; + version?: string; +} + +/** GET /v1/snapshot response body. */ +export type SnapshotResponse = AuthCredentialSnapshot; + +/** POST /v1/credential/:id/refresh response body. */ +export interface CredentialRefreshResponse { + entry: AuthCredentialSnapshotEntry; +} + +/** POST /v1/credential/:id/disable request body. */ +export interface CredentialDisableRequest { + cause: string; +} + +/** POST /v1/credential/:id/disable response body. */ +export interface CredentialDisableResponse { + ok: boolean; +} + +/** + * POST /v1/credential request body. The OAuth `refresh` must be the *real* + * refresh token (not the sentinel) — the broker is the canonical writer. + */ +export interface CredentialUploadRequest { + provider: string; + credential: AuthCredential; +} + +/** POST /v1/credential response body — redacted snapshot of the provider's rows after upsert. */ +export interface CredentialUploadResponse { + entries: AuthCredentialSnapshotEntry[]; +} + +/** + * Default bearer-protected route prefix. The broker exposes `/v1/healthz` + * unauthenticated for liveness probes; everything else requires a bearer. + */ +export const AUTH_BROKER_API_PREFIX = "/v1"; + +/** Default port when none is configured. Loopback-only, no external exposure. */ +export const DEFAULT_AUTH_BROKER_BIND = "127.0.0.1:8765"; + +/** Default broker→provider refresh skew. Refresh credentials this close to expiry. */ +export const DEFAULT_REFRESH_SKEW_MS = 5 * 60_000; + +/** Default broker refresh-loop cadence. */ +export const DEFAULT_REFRESH_INTERVAL_MS = 60_000; diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index a05abaa1d..c5e69cb6e 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -3,9 +3,9 @@ * Handles loading, saving, refreshing credentials, and usage tracking. * * This module defines: - * - `AuthCredentialStore` interface: abstracting persistence (SQLite, memory, etc.) + * - `AuthCredentialStore` interface: persistence abstraction (SQLite, remote vault, …) * - `AuthStorage` class: credential management with round-robin, usage limits, OAuth refresh - * - `AuthCredentialStore`: concrete SQLite-backed implementation + * - `SqliteAuthCredentialStore`: concrete SQLite-backed implementation */ import { Database, type Statement } from "bun:sqlite"; import * as fs from "node:fs/promises"; @@ -78,6 +78,69 @@ export interface StoredAuthCredential { disabledCause: string | null; } +// ───────────────────────────────────────────────────────────────────────────── +// Auth Broker Snapshot Types +// ───────────────────────────────────────────────────────────────────────────── + +/** + * Sentinel value placed in OAuth `refresh` fields when a credential is shared + * via {@link AuthStorage.exportSnapshot}. Refresh tokens never leave the broker; + * clients must call back to refresh. + */ +export const REMOTE_REFRESH_SENTINEL = "__remote__" as const; +export type RemoteRefreshSentinel = typeof REMOTE_REFRESH_SENTINEL; + +/** OAuth credential with refresh token replaced by the broker sentinel. */ +export type RemoteOAuthCredential = Omit & { + refresh: RemoteRefreshSentinel; +}; + +/** Discriminated credential payload as published by the broker. */ +export type SnapshotCredential = ApiKeyCredential | RemoteOAuthCredential; + +export interface AuthCredentialSnapshotEntry { + id: number; + provider: string; + credential: SnapshotCredential; + identityKey: string | null; +} + +/** + * Wire-shaped snapshot exported by {@link AuthStorage.exportSnapshot} and + * served by the auth-broker server on `GET /v1/snapshot`. + */ +export interface AuthCredentialSnapshot { + generatedAt: number; + credentials: AuthCredentialSnapshotEntry[]; +} + +// ───────────────────────────────────────────────────────────────────────────── +// AuthCredentialStore interface +// ───────────────────────────────────────────────────────────────────────────── + +/** + * Persistence abstraction consumed by {@link AuthStorage}. + * + * Concrete implementations: + * - {@link SqliteAuthCredentialStore} — local SQLite-backed store (default). + * - `RemoteAuthCredentialStore` from `./auth-broker` — client-side snapshot of + * a remote broker; mutating methods (`replace*`, `upsert*`, `delete*ForProvider`) + * throw because login flows route through the broker, not the client. + */ +export interface AuthCredentialStore { + close(): void; + listAuthCredentials(provider?: string): StoredAuthCredential[]; + updateAuthCredential(id: number, credential: AuthCredential): void; + deleteAuthCredential(id: number, disabledCause: string): void; + tryDisableAuthCredentialIfMatches(id: number, expectedData: string, disabledCause: string): boolean; + replaceAuthCredentialsForProvider(provider: string, credentials: AuthCredential[]): StoredAuthCredential[]; + upsertAuthCredentialForProvider(provider: string, credential: AuthCredential): StoredAuthCredential[]; + deleteAuthCredentialsForProvider(provider: string, disabledCause: string): void; + getCache(key: string): string | null; + setCache(key: string, value: string, expiresAtSec: number): void; + cleanExpiredCache(): void; +} + // ───────────────────────────────────────────────────────────────────────────── // AuthStorage Options // ───────────────────────────────────────────────────────────────────────────── @@ -117,6 +180,30 @@ export type AuthStorageOptions = { * duplicate credentials (uninteresting hygiene). */ onCredentialDisabled?: (event: CredentialDisabledEvent) => void | Promise; + /** + * Override OAuth refresh. When set, `AuthStorage` calls this instead of the + * per-provider local refresh function. Receives the credential id so the + * implementation can address remote credentials. + * + * Must return updated {@link OAuthCredentials} with at least `access` and + * `expires`. `refresh` may be an opaque sentinel (e.g. `"__remote__"`) when + * the actual refresh token never leaves the broker. + */ + refreshOAuthCredential?: ( + provider: Provider, + credentialId: number, + credential: OAuthCredential, + ) => Promise; + /** + * Human-readable description of the credential store backing this + * AuthStorage instance. Surfaced through {@link AuthStorage.describeCredentialSource} + * so the TUI can show where a token came from (broker URL or local SQLite path). + * + * Examples: + * - `"local ~/.omp/agent/agent.db"` + * - `"broker http://can.internal:8765"` + */ + sourceLabel?: string; }; // ───────────────────────────────────────────────────────────────────────────── @@ -289,6 +376,8 @@ export class AuthStorage { #fallbackResolver?: (provider: string) => string | undefined; #store: AuthCredentialStore; #configValueResolver: (config: string) => Promise; + #refreshOAuthCredentialOverride?: AuthStorageOptions["refreshOAuthCredential"]; + #sourceLabel?: string; #credentialDisabledListeners: Set<(event: CredentialDisabledEvent) => void | Promise> = new Set(); /** * Buffer for credential_disabled events fired while no listener is subscribed. @@ -309,6 +398,8 @@ export class AuthStorage { this.#usageCache = new AuthStorageUsageCache(this.#store); this.#usageFetch = options.usageFetch ?? fetch; this.#usageRequestTimeoutMs = options.usageRequestTimeoutMs ?? DEFAULT_USAGE_REQUEST_TIMEOUT_MS; + this.#refreshOAuthCredentialOverride = options.refreshOAuthCredential; + this.#sourceLabel = options.sourceLabel; if (options.onCredentialDisabled) { // Constructor-registered subscribers are permanent for this AuthStorage's lifetime; // the unsubscribe handle is intentionally discarded. @@ -328,7 +419,7 @@ export class AuthStorage { * @param dbPath - Path to SQLite database */ static async create(dbPath: string, options: AuthStorageOptions = {}): Promise { - const store = await AuthCredentialStore.open(dbPath); + const store = await SqliteAuthCredentialStore.open(dbPath); return new AuthStorage(store, options); } @@ -1264,6 +1355,26 @@ export class AuthStorage { }; } + /** + * Find the stored credential id matching a {@link UsageCredential} so the + * refresh override can address the row. Mirrors the matching logic in + * {@link AuthStorage.#persistRefreshedUsageCredential}. + */ + #findStoredCredentialIdForUsageCredential(provider: Provider, previous: UsageCredential): number | undefined { + const entries = this.#getStoredCredentials(provider); + const match = entries.find(entry => { + if (entry.credential.type !== "oauth") return false; + if (previous.refreshToken && entry.credential.refresh === previous.refreshToken) return true; + if (previous.accessToken && entry.credential.access === previous.accessToken) return true; + return ( + entry.credential.accountId === previous.accountId && + entry.credential.email === previous.email && + entry.credential.projectId === previous.projectId + ); + }); + return match?.id; + } + #persistRefreshedUsageCredential(provider: Provider, previous: UsageCredential, next: UsageCredential): void { const entries = this.#getStoredCredentials(provider); const index = entries.findIndex(entry => { @@ -1312,7 +1423,15 @@ export class AuthStorage { const refreshableCredential = this.#buildRefreshableOauthCredential(request.credential); if (refreshableCredential) { try { - const refreshed = await this.#refreshOAuthCredential(request.provider, refreshableCredential); + const refreshableCredentialId = this.#findStoredCredentialIdForUsageCredential( + request.provider, + request.credential, + ); + const refreshed = await this.#refreshOAuthCredential( + request.provider, + refreshableCredential, + refreshableCredentialId, + ); const refreshedCredential = this.#mergeRefreshedUsageCredential(request.credential, refreshed); this.#persistRefreshedUsageCredential(request.provider, request.credential, refreshedCredential); params = { @@ -1883,9 +2002,11 @@ export class AuthStorage { return; } try { + const credentialId = this.#getStoredCredentials(provider)[candidate.selection.index]?.id; const refreshedCredentials = await this.#refreshOAuthCredential( provider, candidate.selection.credential, + credentialId, ); candidate.selection.credential = { ...candidate.selection.credential, @@ -1927,17 +2048,25 @@ export class AuthStorage { return undefined; } - async #refreshOAuthCredential(provider: Provider, credential: OAuthCredential): Promise { + async #refreshOAuthCredential( + provider: Provider, + credential: OAuthCredential, + credentialId: number | undefined, + ): Promise { if (Date.now() < credential.expires) return credential; - const customProvider = getOAuthProvider(provider); let refreshPromise: Promise; - if (customProvider) { - if (!customProvider.refreshToken) { - throw new Error(`OAuth provider "${provider}" does not support token refresh`); - } - refreshPromise = customProvider.refreshToken(credential); + if (this.#refreshOAuthCredentialOverride && credentialId !== undefined) { + refreshPromise = this.#refreshOAuthCredentialOverride(provider, credentialId, credential); } else { - refreshPromise = refreshOAuthToken(provider as OAuthProvider, credential); + const customProvider = getOAuthProvider(provider); + if (customProvider) { + if (!customProvider.refreshToken) { + throw new Error(`OAuth provider "${provider}" does not support token refresh`); + } + refreshPromise = customProvider.refreshToken(credential); + } else { + refreshPromise = refreshOAuthToken(provider as OAuthProvider, credential); + } } // Bound the refresh so a slow/hanging token endpoint cannot stall credential selection. let timeout: NodeJS.Timeout | undefined; @@ -2014,7 +2143,11 @@ export class AuthStorage { let result: { newCredentials: OAuthCredentials; apiKey: string } | null; const customProvider = getOAuthProvider(provider); if (customProvider) { - const refreshedCredentials = await this.#refreshOAuthCredential(provider, selection.credential); + const refreshedCredentials = await this.#refreshOAuthCredential( + provider, + selection.credential, + this.#getStoredCredentials(provider)[selection.index]?.id, + ); const apiKey = customProvider.getApiKey ? customProvider.getApiKey(refreshedCredentials) : refreshedCredentials.access; @@ -2204,10 +2337,171 @@ export class AuthStorage { // Fall back to custom resolver (e.g., models.json custom providers) return this.#fallbackResolver?.(provider) ?? undefined; } + + // ─── Auth Broker integration ──────────────────────────────────────────── + + /** + * Build a redacted snapshot of all loaded credentials for the auth-broker + * wire. OAuth refresh tokens are replaced with {@link REMOTE_REFRESH_SENTINEL} + * so clients never see the actual refresh token. + * + * Callers must {@link AuthStorage.reload} first when serving a stale snapshot + * (the broker server's HTTP handler does this). + */ + exportSnapshot(): AuthCredentialSnapshot { + const entries: AuthCredentialSnapshotEntry[] = []; + for (const [provider, stored] of this.#data) { + for (const entry of stored) { + const credential = entry.credential; + const redacted: SnapshotCredential = + credential.type === "api_key" ? credential : { ...credential, refresh: REMOTE_REFRESH_SENTINEL }; + entries.push({ + id: entry.id, + provider, + credential: redacted, + identityKey: resolveCredentialIdentityKey(provider, credential), + }); + } + } + return { generatedAt: Date.now(), credentials: entries }; + } + + /** + * Force-refresh the OAuth credential with the given id, bypassing the + * not-yet-expired guard. Used by the auth-broker server to honour + * `POST /v1/credential/:id/refresh`. + * + * Returns the redacted snapshot entry for the refreshed row. + * Throws when no OAuth credential with that id is loaded. + */ + async forceRefreshCredentialById(id: number): Promise { + for (const [provider, entries] of this.#data) { + const index = entries.findIndex(entry => entry.id === id); + if (index === -1) continue; + const target = entries[index]; + if (target.credential.type !== "oauth") { + throw new Error(`Credential ${id} is not OAuth (provider=${provider}, type=${target.credential.type})`); + } + // Pass a clone with expires=0 so the cached not-yet-expired short-circuit + // in #refreshOAuthCredential doesn't suppress the requested refresh. + const stale: OAuthCredential = { ...target.credential, expires: 0 }; + const refreshed = await this.#refreshOAuthCredential(provider as Provider, stale, id); + const updated: OAuthCredential = { + type: "oauth", + access: refreshed.access, + refresh: refreshed.refresh, + expires: refreshed.expires, + accountId: refreshed.accountId ?? target.credential.accountId, + email: refreshed.email ?? target.credential.email, + projectId: refreshed.projectId ?? target.credential.projectId, + enterpriseUrl: refreshed.enterpriseUrl ?? target.credential.enterpriseUrl, + }; + this.#replaceCredentialAt(provider, index, updated); + return { + id, + provider, + credential: { ...updated, refresh: REMOTE_REFRESH_SENTINEL }, + identityKey: resolveCredentialIdentityKey(provider, updated), + }; + } + throw new Error(`No credential with id=${id}`); + } + + /** + * Disable the credential with the given id and emit a + * {@link CredentialDisabledEvent}. Used by the auth-broker server to honour + * `POST /v1/credential/:id/disable`. Returns `false` when no such row exists. + */ + disableCredentialById(id: number, disabledCause: string): boolean { + for (const [provider, entries] of this.#data) { + const index = entries.findIndex(entry => entry.id === id); + if (index === -1) continue; + this.#store.deleteAuthCredential(id, disabledCause); + const next = entries.filter((_value, idx) => idx !== index); + this.#setStoredCredentials(provider, next); + this.#resetProviderAssignments(provider); + this.#emitCredentialDisabled({ provider, disabledCause }); + return true; + } + return false; + } + + /** + * Upsert a credential into the underlying store, refresh the in-memory + * snapshot, and return the redacted snapshot entries for the provider. + * + * Used by the auth-broker server to honour `POST /v1/credential`. The + * persistence layer (`SqliteAuthCredentialStore.upsertAuthCredentialForProvider`) + * does identity-key matching, so re-uploading the same email/account replaces + * the existing row instead of inserting a duplicate. + */ + upsertCredential(provider: string, credential: AuthCredential): AuthCredentialSnapshotEntry[] { + const stored = this.#store.upsertAuthCredentialForProvider(provider, credential); + this.#setStoredCredentials( + provider, + stored.map(entry => ({ id: entry.id, credential: entry.credential })), + ); + this.#resetProviderAssignments(provider); + return stored.map(entry => { + const persisted = entry.credential; + const redacted: SnapshotCredential = + persisted.type === "api_key" ? persisted : { ...persisted, refresh: REMOTE_REFRESH_SENTINEL }; + return { + id: entry.id, + provider: entry.provider, + credential: redacted, + identityKey: resolveCredentialIdentityKey(provider, persisted), + }; + }); + } + + /** + * Describe where the active credential for a provider came from. + * + * Surfaces three layers, highest precedence first: + * 1. Runtime override (`--api-key`). + * 2. Stored credential (the one this session is currently sticky to, or the + * one round-robin would pick next when no session id is supplied). + * 3. Env var / fallback resolver — when no stored credential exists. + * + * The string is purely informational; consumers must not parse it. + */ + describeCredentialSource(provider: string, sessionId?: string): string | undefined { + if (this.#runtimeOverrides.has(provider)) { + return "runtime override (--api-key)"; + } + + const baseLabel = this.#sourceLabel ?? "local store"; + const stored = this.#getStoredCredentials(provider); + if (stored.length === 0) { + if (getEnvApiKey(provider)) return `env ${baseLabel ? `(fallback over ${baseLabel})` : ""}`.trim(); + if (this.#fallbackResolver?.(provider) !== undefined) return `fallback resolver`; + return undefined; + } + + const session = sessionId ? this.#sessionLastCredential.get(provider)?.get(sessionId) : undefined; + // Same selection logic as #selectCredentialByType for "no session" lookups: prefer + // the type with stored credentials, lean OAuth before api_key. We don't run the + // full round-robin here because describing the source shouldn't advance the index. + const preferredType: AuthCredential["type"] = + session?.type ?? (stored.some(entry => entry.credential.type === "oauth") ? "oauth" : "api_key"); + const typed = stored + .map((entry, index) => ({ entry, index })) + .filter(({ entry }) => entry.credential.type === preferredType); + if (typed.length === 0) return baseLabel; + const index = session?.index ?? typed[0].index; + const chosen = stored[index] ?? typed[0].entry; + const credential = chosen.credential; + const identity = + credential.type === "oauth" + ? (credential.email ?? credential.accountId ?? credential.projectId ?? `cred ${chosen.id}`) + : `cred ${chosen.id}`; + return `${baseLabel} · ${preferredType} #${chosen.id} (${identity})`; + } } // ───────────────────────────────────────────────────────────────────────────── -// AuthCredentialStore +// SqliteAuthCredentialStore // ───────────────────────────────────────────────────────────────────────────── /** Row shape for auth_credentials table queries */ @@ -2389,11 +2683,14 @@ function extractOAuthTokenIdentifiers(token: string | undefined): string[] | und } } /** - * Standalone SQLite-backed implementation of AuthCredentialStore interface. - * Used by the pi-ai CLI and as the default store for AuthStorage.create(). - * Also has convenience methods for simple CRUD (saveOAuth, getOAuth, etc.). + * Default SQLite-backed implementation of {@link AuthCredentialStore}. + * + * Used by the pi-ai CLI and as the default store for `AuthStorage.create()`. + * Also exposes convenience methods (`saveOAuth`, `getOAuth`, `saveApiKey`, + * `getApiKey`, `listProviders`, `deleteProvider`) that callers can use directly + * without going through `AuthStorage`. */ -export class AuthCredentialStore { +export class SqliteAuthCredentialStore implements AuthCredentialStore { #db: Database; #listActiveStmt: Statement; #listActiveByProviderStmt: Statement; @@ -2447,7 +2744,7 @@ export class AuthCredentialStore { this.#deleteExpiredCacheStmt = this.#db.prepare(`DELETE FROM cache WHERE expires_at <= ${SQLITE_NOW_EPOCH}`); } - static async open(dbPath: string = getAgentDbPath()): Promise { + static async open(dbPath: string = getAgentDbPath()): Promise { const dir = path.dirname(dbPath); const dirExists = await fs .stat(dir) @@ -2464,7 +2761,7 @@ export class AuthCredentialStore { // Ignore chmod failures (e.g., Windows) } - return new AuthCredentialStore(db); + return new SqliteAuthCredentialStore(db); } #initializeSchema(): void { @@ -2493,7 +2790,7 @@ export class AuthCredentialStore { const schemaVersion = this.#readAuthSchemaVersion() ?? this.#inferAuthSchemaVersion(); const shouldWriteSchemaVersion = schemaVersion <= AUTH_SCHEMA_VERSION; if (schemaVersion > AUTH_SCHEMA_VERSION) { - logger.warn("AuthCredentialStore schema version mismatch", { + logger.warn("SqliteAuthCredentialStore schema version mismatch", { current: schemaVersion, expected: AUTH_SCHEMA_VERSION, }); diff --git a/packages/ai/src/cli.ts b/packages/ai/src/cli.ts index 0f544ea78..78c1940be 100755 --- a/packages/ai/src/cli.ts +++ b/packages/ai/src/cli.ts @@ -1,6 +1,6 @@ #!/usr/bin/env bun import * as readline from "node:readline"; -import { AuthCredentialStore } from "./auth-storage"; +import { SqliteAuthCredentialStore } from "./auth-storage"; import { getOAuthProviders } from "./utils/oauth"; import type { OAuthCredentials, OAuthProvider } from "./utils/oauth/types"; @@ -60,7 +60,7 @@ async function login(provider: OAuthProvider): Promise { const rl = readline.createInterface({ input: process.stdin, output: process.stdout }); const promptFn = (msg: string) => prompt(rl, `${msg} `); - const storage = await AuthCredentialStore.open(); + const storage = await SqliteAuthCredentialStore.open(); try { let credentials: OAuthCredentials; @@ -387,7 +387,7 @@ Examples: } if (command === "status") { - const storage = await AuthCredentialStore.open(); + const storage = await SqliteAuthCredentialStore.open(); try { const providers = storage.listProviders(); if (providers.length === 0) { @@ -426,7 +426,7 @@ Examples: if (command === "logout") { let provider = args[1] as OAuthProvider | undefined; - const storage = await AuthCredentialStore.open(); + const storage = await SqliteAuthCredentialStore.open(); try { if (!provider) { diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 89088192a..e007493cd 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -1,5 +1,6 @@ export { type ZodType, z } from "zod/v4"; export * from "./api-registry"; +export * from "./auth-broker"; export * from "./auth-storage"; export * from "./model-cache"; export * from "./model-manager"; @@ -38,6 +39,13 @@ export * from "./utils/anthropic-auth"; export * from "./utils/discovery"; export * from "./utils/event-stream"; export * from "./utils/h2-fetch"; +export * from "./utils/oauth"; +export type { + OAuthCredentials, + OAuthProvider, + OAuthProviderId, + OAuthProviderInfo, +} from "./utils/oauth/types"; export * from "./utils/overflow"; export * from "./utils/retry"; export * from "./utils/schema"; diff --git a/packages/ai/src/provider-details.ts b/packages/ai/src/provider-details.ts index d775d091f..40a54923f 100644 --- a/packages/ai/src/provider-details.ts +++ b/packages/ai/src/provider-details.ts @@ -16,6 +16,12 @@ export interface ProviderDetailsContext { model: Model; sessionId?: string; authMode?: string; + /** + * Human-readable description of the active credential, e.g. + * `"broker http://can.internal:8765 · oauth #5 (foo@bar.com)"`. + * Rendered as a `Source` field; omitted when undefined. + */ + credentialSource?: string; preferWebsockets?: boolean; providerSessionState?: Map; } @@ -28,6 +34,9 @@ export function getProviderDetails(context: ProviderDetailsContext): ProviderDet { label: "Auth", value: context.authMode ?? "auto" }, { label: "Endpoint", value: endpoint }, ]; + if (context.credentialSource) { + fields.push({ label: "Source", value: context.credentialSource }); + } if (context.model.api === "openai-codex-responses") { const codexDetails = getOpenAICodexTransportDetails(context.model as Model<"openai-codex-responses">, { diff --git a/packages/ai/src/utils/anthropic-auth.ts b/packages/ai/src/utils/anthropic-auth.ts index dbfddcc87..a00b20a50 100644 --- a/packages/ai/src/utils/anthropic-auth.ts +++ b/packages/ai/src/utils/anthropic-auth.ts @@ -9,7 +9,7 @@ * 5. Generic Anthropic fallback (ANTHROPIC_API_KEY / ANTHROPIC_BASE_URL) */ import { $env, getAgentDbPath } from "@oh-my-pi/pi-utils"; -import { type AuthCredential, AuthCredentialStore } from "../auth-storage"; +import { type AuthCredential, type AuthCredentialStore, SqliteAuthCredentialStore } from "../auth-storage"; import { buildAnthropicHeaders as buildProviderAnthropicHeaders, normalizeAnthropicBaseUrl, @@ -80,7 +80,7 @@ function toAnthropicOAuthCredential(credential: AuthCredential): AnthropicOAuthC */ async function readAnthropicOAuthCredentials(store?: AuthCredentialStore): Promise { const ownsStore = !store; - const effectiveStore = store ?? (await AuthCredentialStore.open(getAgentDbPath())); + const effectiveStore = store ?? (await SqliteAuthCredentialStore.open(getAgentDbPath())); try { const records = effectiveStore.listAuthCredentials("anthropic"); const credentials: AnthropicOAuthCredential[] = []; @@ -133,7 +133,7 @@ export async function findAnthropicAuth(store?: AuthCredentialStore): Promise record.credential.type === "api_key"); + if (apiKeyRecord && apiKeyRecord.credential.type === "api_key") { return { - apiKey: storedApiKey, + apiKey: apiKeyRecord.credential.key, baseUrl: resolveAnthropicBaseUrlFromEnv() ?? DEFAULT_BASE_URL, - isOAuth: isOAuthToken(storedApiKey), + isOAuth: isOAuthToken(apiKeyRecord.credential.key), }; } } finally { diff --git a/packages/ai/test/anthropic-oauth.test.ts b/packages/ai/test/anthropic-oauth.test.ts index 3a66ab1d6..66345cb4d 100644 --- a/packages/ai/test/anthropic-oauth.test.ts +++ b/packages/ai/test/anthropic-oauth.test.ts @@ -2,7 +2,7 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore } from "../src/auth-storage"; +import { SqliteAuthCredentialStore } from "../src/auth-storage"; import { buildAnthropicUrl, findAnthropicAuth } from "../src/utils/anthropic-auth"; import { AnthropicOAuthFlow, refreshAnthropicToken } from "../src/utils/oauth/anthropic"; import { withEnv } from "./helpers"; @@ -199,7 +199,7 @@ describe("anthropic auth resolution", () => { const tmpDir = path.join(os.tmpdir(), `pi-ai-auth-${Date.now()}-${Math.random().toString(16).slice(2)}`); fs.mkdirSync(tmpDir, { recursive: true }); const dbPath = path.join(tmpDir, "agent.db"); - const store = await AuthCredentialStore.open(dbPath); + const store = await SqliteAuthCredentialStore.open(dbPath); try { store.replaceAuthCredentialsForProvider("anthropic", [ { type: "oauth", access: "sk-ant-oat-db", refresh: "refresh", expires: Date.now() + 20 * 60 * 1000 }, @@ -231,7 +231,7 @@ describe("anthropic auth resolution", () => { const tmpDir = path.join(os.tmpdir(), `pi-ai-auth-${Date.now()}-${Math.random().toString(16).slice(2)}`); fs.mkdirSync(tmpDir, { recursive: true }); const dbPath = path.join(tmpDir, "agent.db"); - const store = await AuthCredentialStore.open(dbPath); + const store = await SqliteAuthCredentialStore.open(dbPath); try { store.replaceAuthCredentialsForProvider("anthropic", [ { type: "oauth", access: "sk-ant-oat-db", refresh: "refresh", expires: Date.now() + 20 * 60 * 1000 }, @@ -261,7 +261,7 @@ describe("anthropic auth resolution", () => { const tmpDir = path.join(os.tmpdir(), `pi-ai-auth-${Date.now()}-${Math.random().toString(16).slice(2)}`); fs.mkdirSync(tmpDir, { recursive: true }); const dbPath = path.join(tmpDir, "agent.db"); - const store = await AuthCredentialStore.open(dbPath); + const store = await SqliteAuthCredentialStore.open(dbPath); try { store.replaceAuthCredentialsForProvider("anthropic", [{ type: "api_key", key: "sk-ant-api-db" }]); await withEnv( diff --git a/packages/ai/test/auth-broker-refresher.test.ts b/packages/ai/test/auth-broker-refresher.test.ts new file mode 100644 index 000000000..52e192963 --- /dev/null +++ b/packages/ai/test/auth-broker-refresher.test.ts @@ -0,0 +1,150 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { AuthBrokerRefresher, AuthStorage, SqliteAuthCredentialStore } from "../src"; +import * as oauthUtils from "../src/utils/oauth"; + +const ANTHROPIC_ENV = ["ANTHROPIC_API_KEY", "ANTHROPIC_OAUTH_TOKEN"] as const; +const savedEnv: Partial> = {}; + +describe("AuthBrokerRefresher", () => { + let tempDir = ""; + let store: SqliteAuthCredentialStore | undefined; + let storage: AuthStorage | undefined; + + beforeEach(async () => { + for (const key of ANTHROPIC_ENV) { + savedEnv[key] = process.env[key]; + delete process.env[key]; + } + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "auth-broker-refresher-")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + storage?.close(); + store?.close(); + await fs.rm(tempDir, { recursive: true, force: true }); + for (const key of ANTHROPIC_ENV) { + if (savedEnv[key] === undefined) delete process.env[key]; + else process.env[key] = savedEnv[key]; + } + }); + + test("refreshes credentials inside the skew window", async () => { + const now = 1_700_000_000_000; + const skew = 5 * 60_000; + // Credential expires in 1 minute — well within the 5-min skew → must refresh. + store!.saveOAuth("anthropic", { + access: "old", + refresh: "old-refresh", + expires: now + 60_000, + accountId: "a", + }); + const refreshSpy = vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue({ + access: "fresh", + refresh: "fresh-refresh", + expires: now + 2 * 60 * 60_000, + accountId: "a", + }); + + storage = new AuthStorage(store!); + await storage.reload(); + const refresher = new AuthBrokerRefresher({ + storage, + refreshSkewMs: skew, + now: () => now, + }); + await refresher.tick(); + + expect(refreshSpy).toHaveBeenCalledTimes(1); + const persisted = store!.getOAuth("anthropic"); + expect(persisted?.access).toBe("fresh"); + expect(persisted?.refresh).toBe("fresh-refresh"); + }); + + test("does not refresh credentials safely outside the skew window", async () => { + const now = 1_700_000_000_000; + const skew = 5 * 60_000; + store!.saveOAuth("anthropic", { + access: "ok", + refresh: "ok-refresh", + expires: now + 60 * 60_000, // 1 hour out + accountId: "a", + }); + const refreshSpy = vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue({ + access: "should-not-run", + refresh: "x", + expires: now, + }); + + storage = new AuthStorage(store!); + await storage.reload(); + const refresher = new AuthBrokerRefresher({ + storage, + refreshSkewMs: skew, + now: () => now, + }); + await refresher.tick(); + + expect(refreshSpy).not.toHaveBeenCalled(); + }); + + test("disables credentials on definitive failure (invalid_grant)", async () => { + const now = 1_700_000_000_000; + store!.saveOAuth("anthropic", { + access: "old", + refresh: "old-refresh", + expires: now + 60_000, + accountId: "a", + }); + vi.spyOn(oauthUtils, "refreshOAuthToken").mockRejectedValue(new Error("invalid_grant")); + + storage = new AuthStorage(store!); + const disableEvents: string[] = []; + storage.onCredentialDisabled(event => { + disableEvents.push(event.disabledCause); + }); + await storage.reload(); + const refresher = new AuthBrokerRefresher({ + storage, + refreshSkewMs: 5 * 60_000, + now: () => now, + }); + await refresher.tick(); + + expect(disableEvents).toHaveLength(1); + expect(disableEvents[0]).toMatch(/invalid_grant/); + // The active row is now disabled; storage.exportSnapshot reflects it. + expect(storage.exportSnapshot().credentials).toHaveLength(0); + }); + + test("keeps credentials on transient failures (timeout/network)", async () => { + const now = 1_700_000_000_000; + store!.saveOAuth("anthropic", { + access: "old", + refresh: "old-refresh", + expires: now + 60_000, + accountId: "a", + }); + vi.spyOn(oauthUtils, "refreshOAuthToken").mockRejectedValue(new Error("fetch failed: ECONNREFUSED")); + + storage = new AuthStorage(store!); + const disableEvents: string[] = []; + storage.onCredentialDisabled(event => { + disableEvents.push(event.disabledCause); + }); + await storage.reload(); + const refresher = new AuthBrokerRefresher({ + storage, + refreshSkewMs: 5 * 60_000, + now: () => now, + }); + await refresher.tick(); + + expect(disableEvents).toHaveLength(0); + expect(storage.exportSnapshot().credentials).toHaveLength(1); + }); +}); diff --git a/packages/ai/test/auth-broker-wire.test.ts b/packages/ai/test/auth-broker-wire.test.ts new file mode 100644 index 000000000..d3d3d57c1 --- /dev/null +++ b/packages/ai/test/auth-broker-wire.test.ts @@ -0,0 +1,138 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { + AuthBrokerClient, + type AuthBrokerServerHandle, + AuthStorage, + REMOTE_REFRESH_SENTINEL, + SqliteAuthCredentialStore, + startAuthBroker, +} from "../src"; +import * as oauthUtils from "../src/utils/oauth"; + +const ANTHROPIC_ENV = ["ANTHROPIC_API_KEY", "ANTHROPIC_OAUTH_TOKEN"] as const; +const savedEnv: Partial> = {}; + +function mintOAuthCredential(suffix: string, expires: number) { + return { + type: "oauth" as const, + access: `access-${suffix}`, + refresh: `refresh-${suffix}`, + expires, + accountId: `account-${suffix}`, + email: `${suffix}@example.com`, + }; +} + +describe("auth-broker wire surface", () => { + let tempDir = ""; + let store: SqliteAuthCredentialStore | undefined; + let storage: AuthStorage | undefined; + let handle: AuthBrokerServerHandle | undefined; + let token = ""; + + beforeEach(async () => { + for (const key of ANTHROPIC_ENV) { + savedEnv[key] = process.env[key]; + delete process.env[key]; + } + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "auth-broker-wire-")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + store.saveOAuth("anthropic", mintOAuthCredential("a", Date.now() + 60_000)); + storage = new AuthStorage(store); + await storage.reload(); + token = "test-bearer"; + handle = startAuthBroker({ + storage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + await handle?.close(); + storage?.close(); + store?.close(); + await fs.rm(tempDir, { recursive: true, force: true }); + for (const key of ANTHROPIC_ENV) { + if (savedEnv[key] === undefined) delete process.env[key]; + else process.env[key] = savedEnv[key]; + } + }); + + test("GET /v1/healthz returns ok without auth", async () => { + const res = await fetch(`${handle!.url}/v1/healthz`); + expect(res.status).toBe(200); + const body = (await res.json()) as { ok: boolean }; + expect(body.ok).toBe(true); + }); + + test("GET /v1/snapshot requires bearer and redacts refresh tokens", async () => { + const unauthorized = await fetch(`${handle!.url}/v1/snapshot`); + expect(unauthorized.status).toBe(401); + + const client = new AuthBrokerClient({ url: handle!.url, token }); + const snapshot = await client.fetchSnapshot(); + expect(snapshot.credentials).toHaveLength(1); + const entry = snapshot.credentials[0]; + expect(entry.provider).toBe("anthropic"); + expect(entry.credential.type).toBe("oauth"); + if (entry.credential.type === "oauth") { + expect(entry.credential.access).toBe("access-a"); + // Refresh token is replaced with the wire sentinel — clients never see it. + expect(entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL); + } + }); + + test("POST /v1/credential/:id/refresh forces a refresh and persists the new credential", async () => { + const refreshed = { + access: "access-rotated", + refresh: "refresh-rotated", + expires: Date.now() + 120_000, + accountId: "account-a", + email: "a@example.com", + }; + vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue(refreshed); + + const initialSnapshot = await new AuthBrokerClient({ url: handle!.url, token }).fetchSnapshot(); + const id = initialSnapshot.credentials[0].id; + + const client = new AuthBrokerClient({ url: handle!.url, token }); + const result = await client.refreshCredential(id); + expect(result.entry.id).toBe(id); + if (result.entry.credential.type === "oauth") { + expect(result.entry.credential.access).toBe("access-rotated"); + expect(result.entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL); + } + + // Underlying SQLite row was updated with the *real* refresh token (no sentinel). + const persisted = store!.getOAuth("anthropic"); + expect(persisted?.access).toBe("access-rotated"); + expect(persisted?.refresh).toBe("refresh-rotated"); + }); + + test("POST /v1/credential/:id/disable soft-deletes the credential and surfaces 404 thereafter", async () => { + const client = new AuthBrokerClient({ url: handle!.url, token }); + const initialSnapshot = await client.fetchSnapshot(); + const id = initialSnapshot.credentials[0].id; + + const result = await client.disableCredential(id, "revoked by user"); + expect(result.ok).toBe(true); + + const after = await client.fetchSnapshot(); + expect(after.credentials).toHaveLength(0); + + await expect(client.refreshCredential(id)).rejects.toThrow(); + }); + + test("Unknown route returns 404", async () => { + const res = await fetch(`${handle!.url}/v1/nope`, { + headers: { Authorization: `Bearer ${token}` }, + }); + expect(res.status).toBe(404); + }); +}); diff --git a/packages/ai/test/auth-storage-api-key-login.test.ts b/packages/ai/test/auth-storage-api-key-login.test.ts index 48abc62a8..62e3a5a7d 100644 --- a/packages/ai/test/auth-storage-api-key-login.test.ts +++ b/packages/ai/test/auth-storage-api-key-login.test.ts @@ -4,7 +4,7 @@ import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore, AuthStorage } from "../src/auth-storage"; +import { AuthStorage, SqliteAuthCredentialStore } from "../src/auth-storage"; import * as kagiModule from "../src/utils/oauth/kagi"; import * as ollamaCloudModule from "../src/utils/oauth/ollama-cloud"; @@ -23,7 +23,7 @@ function countCredentialRows(dbPath: string, provider: string): number { describe("AuthStorage api-key login replacement", () => { let tempDir = ""; let dbPath = ""; - let store: AuthCredentialStore | null = null; + let store: SqliteAuthCredentialStore | null = null; let authStorage: AuthStorage | null = null; let loginKagiSpy: Mock; let loginOllamaCloudSpy: Mock; @@ -31,7 +31,7 @@ describe("AuthStorage api-key login replacement", () => { beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-api-key-login-")); dbPath = path.join(tempDir, "agent.db"); - store = await AuthCredentialStore.open(dbPath); + store = await SqliteAuthCredentialStore.open(dbPath); authStorage = new AuthStorage(store); loginKagiSpy = vi.spyOn(kagiModule, "loginKagi"); loginOllamaCloudSpy = vi.spyOn(ollamaCloudModule, "loginOllamaCloud"); diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 0a74c4c4d..99bffa787 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -2,7 +2,7 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore, AuthStorage } from "../src/auth-storage"; +import { type AuthCredentialStore, AuthStorage, SqliteAuthCredentialStore } from "../src/auth-storage"; import type { UsageLimit, UsageProvider, UsageReport } from "../src/usage"; import * as oauthUtils from "../src/utils/oauth"; import type { OAuthCredentials } from "../src/utils/oauth/types"; @@ -120,7 +120,7 @@ describe("AuthStorage codex oauth ranking", () => { beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-codex-selection-")); - store = await AuthCredentialStore.open(path.join(tempDir, "agent.db")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); authStorage = new AuthStorage(store, { usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined), }); @@ -590,7 +590,7 @@ describe("AuthStorage claude oauth ranking", () => { beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-claude-selection-")); - store = await AuthCredentialStore.open(path.join(tempDir, "agent.db")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); authStorage = new AuthStorage(store, { usageProviderResolver: provider => (provider === "anthropic" ? usageProvider : undefined), }); diff --git a/packages/ai/test/auth-storage-credential-disabled-event.test.ts b/packages/ai/test/auth-storage-credential-disabled-event.test.ts index 6c3bb0f6d..01184eb53 100644 --- a/packages/ai/test/auth-storage-credential-disabled-event.test.ts +++ b/packages/ai/test/auth-storage-credential-disabled-event.test.ts @@ -2,7 +2,12 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore, AuthStorage, type CredentialDisabledEvent } from "../src/auth-storage"; +import { + type AuthCredentialStore, + AuthStorage, + type CredentialDisabledEvent, + SqliteAuthCredentialStore, +} from "../src/auth-storage"; import * as oauthUtils from "../src/utils/oauth"; // Env vars short-circuit AuthStorage.getApiKey before the OAuth refresh path runs; suppress @@ -29,7 +34,7 @@ describe("AuthStorage credential_disabled subscriptions", () => { const stores: AuthCredentialStore[] = []; const openStorage = async (options?: ConstructorParameters[1]): Promise => { - const store = await AuthCredentialStore.open(path.join(tempDir, `agent-${stores.length}.db`)); + const store = await SqliteAuthCredentialStore.open(path.join(tempDir, `agent-${stores.length}.db`)); stores.push(store); return new AuthStorage(store, options); }; diff --git a/packages/ai/test/auth-storage-email-dedupe.test.ts b/packages/ai/test/auth-storage-email-dedupe.test.ts index e0cd8ccf7..434d703aa 100644 --- a/packages/ai/test/auth-storage-email-dedupe.test.ts +++ b/packages/ai/test/auth-storage-email-dedupe.test.ts @@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore, AuthStorage, type OAuthCredential } from "../src/auth-storage"; +import { AuthStorage, type OAuthCredential, SqliteAuthCredentialStore } from "../src/auth-storage"; const LEGACY_TIMESTAMP = 1_700_000_000; @@ -105,13 +105,13 @@ function readTableSql(dbPath: string, tableName: string): string | null { describe("AuthStorage openai-codex email dedupe", () => { let tempDir = ""; let dbPath = ""; - let store: AuthCredentialStore | null = null; + let store: SqliteAuthCredentialStore | null = null; let authStorage: AuthStorage | null = null; beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-email-dedupe-")); dbPath = path.join(tempDir, "agent.db"); - store = await AuthCredentialStore.open(dbPath); + store = await SqliteAuthCredentialStore.open(dbPath); authStorage = new AuthStorage(store); }); @@ -255,8 +255,8 @@ describe("AuthStorage openai-codex email dedupe", () => { it("saveOAuth does not delete accounts missing from stale AuthStorage cache", async () => { if (!store || !dbPath) throw new Error("test setup failed"); - const staleStore = await AuthCredentialStore.open(dbPath); - const freshStore = await AuthCredentialStore.open(dbPath); + const staleStore = await SqliteAuthCredentialStore.open(dbPath); + const freshStore = await SqliteAuthCredentialStore.open(dbPath); const staleAuthStorage = new AuthStorage(staleStore); try { staleStore.saveOAuth( @@ -396,7 +396,7 @@ describe("AuthStorage openai-codex email dedupe", () => { ); legacyDb.close(); - const migratedStore = await AuthCredentialStore.open(legacyDbPath); + const migratedStore = await SqliteAuthCredentialStore.open(legacyDbPath); try { expect(readStoredIdentityRows(legacyDbPath, "anthropic")).toEqual([ { identity_key: "email:legacy-anthropic@example.com", disabled_cause: null }, @@ -427,7 +427,7 @@ describe("AuthStorage openai-codex email dedupe", () => { if (!tempDir) throw new Error("test setup failed"); const freshDbPath = path.join(tempDir, "fresh-schema-agent.db"); - const freshStore = await AuthCredentialStore.open(freshDbPath); + const freshStore = await SqliteAuthCredentialStore.open(freshDbPath); try { expect(readAuthSchemaVersion(freshDbPath)).toBe(4); expect(readTableSql(freshDbPath, "auth_credentials")).not.toContain("unixepoch("); @@ -461,7 +461,7 @@ describe("AuthStorage openai-codex email dedupe", () => { `); futureDb.close(); - const reopenedStore = await AuthCredentialStore.open(futureDbPath); + const reopenedStore = await SqliteAuthCredentialStore.open(futureDbPath); try { expect(readAuthSchemaVersion(futureDbPath)).toBe(5); } finally { @@ -512,7 +512,7 @@ describe("AuthStorage openai-codex email dedupe", () => { ); legacyDb.close(); - const migratedStore = await AuthCredentialStore.open(legacyDbPath); + const migratedStore = await SqliteAuthCredentialStore.open(legacyDbPath); try { expect(readAuthSchemaVersion(legacyDbPath)).toBe(4); expect(readTableSql(legacyDbPath, "auth_credentials")).not.toContain("unixepoch("); @@ -566,7 +566,7 @@ describe("AuthStorage openai-codex email dedupe", () => { ); legacyDb.close(); - const migratedStore = await AuthCredentialStore.open(legacyDbPath); + const migratedStore = await SqliteAuthCredentialStore.open(legacyDbPath); try { expect(readStoredIdentityRows(legacyDbPath, "openai-codex")).toEqual([ { identity_key: "email:legacy-v1@example.com", disabled_cause: null }, @@ -608,7 +608,7 @@ describe("AuthStorage openai-codex email dedupe", () => { ); legacyDb.close(); - const migratedStore = await AuthCredentialStore.open(legacyDbPath); + const migratedStore = await SqliteAuthCredentialStore.open(legacyDbPath); try { expect(migratedStore.listAuthCredentials("openai-codex")).toHaveLength(0); expect(readStoredIdentityRows(legacyDbPath, "openai-codex")).toEqual([ diff --git a/packages/ai/test/auth-storage-oauth-refresh-race.test.ts b/packages/ai/test/auth-storage-oauth-refresh-race.test.ts index 78eae8a55..59864f631 100644 --- a/packages/ai/test/auth-storage-oauth-refresh-race.test.ts +++ b/packages/ai/test/auth-storage-oauth-refresh-race.test.ts @@ -2,7 +2,12 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { AuthCredentialStore, AuthStorage, type CredentialDisabledEvent } from "../src/auth-storage"; +import { + type AuthCredentialStore, + AuthStorage, + type CredentialDisabledEvent, + SqliteAuthCredentialStore, +} from "../src/auth-storage"; import * as oauthUtils from "../src/utils/oauth"; import { withEnv } from "./helpers"; @@ -19,7 +24,7 @@ describe("AuthStorage OAuth refresh race", () => { beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-oauth-race-")); - store = await AuthCredentialStore.open(path.join(tempDir, "agent.db")); + store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); events = []; authStorage = new AuthStorage(store, { onCredentialDisabled: event => { diff --git a/packages/ai/test/remote-auth-store.test.ts b/packages/ai/test/remote-auth-store.test.ts new file mode 100644 index 000000000..0929c72d9 --- /dev/null +++ b/packages/ai/test/remote-auth-store.test.ts @@ -0,0 +1,118 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { + AuthBrokerClient, + type AuthBrokerServerHandle, + AuthStorage, + REMOTE_REFRESH_SENTINEL, + RemoteAuthCredentialStore, + SqliteAuthCredentialStore, + startAuthBroker, +} from "../src"; +import * as oauthUtils from "../src/utils/oauth"; + +const ANTHROPIC_ENV = ["ANTHROPIC_API_KEY", "ANTHROPIC_OAUTH_TOKEN"] as const; +const savedEnv: Partial> = {}; + +describe("RemoteAuthCredentialStore + AuthStorage integration", () => { + let tempDir = ""; + let serverStore: SqliteAuthCredentialStore | undefined; + let serverStorage: AuthStorage | undefined; + let handle: AuthBrokerServerHandle | undefined; + const token = "remote-bearer"; + + beforeEach(async () => { + for (const key of ANTHROPIC_ENV) { + savedEnv[key] = process.env[key]; + delete process.env[key]; + } + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "auth-broker-remote-")); + serverStore = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + serverStore.saveOAuth("anthropic", { + access: "server-access-1", + refresh: "server-refresh-1", + expires: Date.now() - 60_000, // expired so refresh is forced + accountId: "account-1", + email: "a@example.com", + }); + serverStorage = new AuthStorage(serverStore); + await serverStorage.reload(); + handle = startAuthBroker({ + storage: serverStorage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + await handle?.close(); + serverStorage?.close(); + serverStore?.close(); + await fs.rm(tempDir, { recursive: true, force: true }); + for (const key of ANTHROPIC_ENV) { + if (savedEnv[key] === undefined) delete process.env[key]; + else process.env[key] = savedEnv[key]; + } + }); + + test("client-side AuthStorage refreshes via broker override, never via local OAuth path", async () => { + // Real refresh executed by the broker server; mock surfaces the rotated tokens. + const rotated = { + access: "server-access-rotated", + refresh: "server-refresh-rotated", + expires: Date.now() + 120_000, + accountId: "account-1", + email: "a@example.com", + }; + const refreshSpy = vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue(rotated); + + const brokerClient = new AuthBrokerClient({ url: handle!.url, token }); + const initialSnapshot = await brokerClient.fetchSnapshot(); + expect(initialSnapshot.credentials).toHaveLength(1); + + const remoteStore = new RemoteAuthCredentialStore({ + client: brokerClient, + initialSnapshot, + }); + + let overrideCalls = 0; + const clientStorage = new AuthStorage(remoteStore, { + refreshOAuthCredential: async (_provider, credentialId, _credential) => { + overrideCalls += 1; + const { entry } = await brokerClient.refreshCredential(credentialId); + if (entry.credential.type !== "oauth") throw new Error("unexpected"); + return { + access: entry.credential.access, + refresh: REMOTE_REFRESH_SENTINEL, + expires: entry.credential.expires, + accountId: entry.credential.accountId, + email: entry.credential.email, + }; + }, + }); + await clientStorage.reload(); + + const apiKey = await clientStorage.getApiKey("anthropic"); + expect(apiKey).toBe("server-access-rotated"); + expect(overrideCalls).toBe(1); + // The local oauth refresh helper was used exactly once — by the broker server. + expect(refreshSpy).toHaveBeenCalledTimes(1); + clientStorage.close(); + }); + + test("RemoteAuthCredentialStore rejects writes from the client", () => { + const remoteStore = new RemoteAuthCredentialStore({ + client: new AuthBrokerClient({ url: handle!.url, token }), + }); + expect(() => remoteStore.replaceAuthCredentialsForProvider("anthropic", [])).toThrow(/read-only/); + expect(() => remoteStore.upsertAuthCredentialForProvider("anthropic", { type: "api_key", key: "x" })).toThrow( + /read-only/, + ); + expect(() => remoteStore.deleteAuthCredentialsForProvider("anthropic", "x")).toThrow(/read-only/); + remoteStore.close(); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f72876a35..589720019 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -7,6 +7,18 @@ - Renamed the embedded-documentation internal URL scheme from `pi://` to `omp://`. `OmpProtocolHandler` replaces `PiProtocolHandler`; update any external references accordingly. - Removed the `StringEnum` re-export from `@oh-my-pi/pi-coding-agent`. Custom tools and extensions should use `z.enum([...])` directly via the injected `pi.zod`. - Replaced the `eval` tool's LARK-grammar `input` string with a structured `cells` array. Each cell is `{ language: "py" | "js", code, title?, timeout?, reset? }`. Removed the implicit/sniffed language path, the `*** Cell` / `*** End` / `*** Abort` markers, and the per-cell `t:` unit suffixes — `timeout` is now seconds (1-600). + +### Added + +- Added `omp auth-broker` subcommand for running and consuming a hosted credential vault. + - `serve [--bind=host:port]` — boots a local broker against the SQLite store at `$AGENT_DB_PATH`. + - `token [--regenerate]` — prints (and rotates) the bearer token stored at `~/.omp/auth-broker.token`. + - `login [--via=user@host] [--dry-run]` — drives the OAuth flow locally or via SSH `-L` tunnel into a remote broker (callback ports pinned per provider). + - `logout ` — disables every credential for the given provider in the local SQLite store. + - `import [--provider=] [--include-disabled] [--dry-run]` — imports CLIProxyAPI-style JSON credential dumps (`~/.cliproxy/auth/*.json`). When `OMP_AUTH_BROKER_URL` is configured, credentials are uploaded to the remote broker via `POST /v1/credential`; otherwise they go into the local SQLite store. JSON `type` is mapped to omp providers (`claude` → `anthropic`, `codex` → `openai-codex`, `gemini[-cli]` → `google-gemini-cli`, `antigravity` → `google-antigravity`); `--provider` overrides the mapping for unrecognized types. + - `status` — pings the configured remote broker (`OMP_AUTH_BROKER_URL`). +- Added remote credential vault support to `discoverAuthStorage`. Configure via env (`OMP_AUTH_BROKER_URL` / `OMP_AUTH_BROKER_TOKEN`) or by setting `auth.broker.url` and `auth.broker.token` in `~/.omp/agent/config.yml` (hidden from the settings UI; supports `!command` resolution). Falls back to `~/.omp/auth-broker.token` when no token is provided inline. Otherwise behavior is unchanged. + ### Changed - Changed TTSR `interruptMode` semantics so a non-interrupting decision on a tool-source match now folds the rule reminder into that specific tool's `toolResult` content instead of queuing a loop-wide deferred follow-up turn. Text/thinking matches keep the previous deferred-injection behavior. diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts index f672269c8..fe077137f 100755 --- a/packages/coding-agent/src/cli.ts +++ b/packages/coding-agent/src/cli.ts @@ -32,6 +32,7 @@ process.title = APP_NAME; const commands: CommandEntry[] = [ { name: "launch", load: () => import("./commands/launch").then(m => m.default) }, { name: "acp", load: () => import("./commands/acp").then(m => m.default) }, + { name: "auth-broker", load: () => import("./commands/auth-broker").then(m => m.default) }, { name: "agents", load: () => import("./commands/agents").then(m => m.default) }, { name: "commit", load: () => import("./commands/commit").then(m => m.default) }, { name: "config", load: () => import("./commands/config").then(m => m.default) }, diff --git a/packages/coding-agent/src/cli/auth-broker-cli.ts b/packages/coding-agent/src/cli/auth-broker-cli.ts new file mode 100644 index 000000000..a354d64e8 --- /dev/null +++ b/packages/coding-agent/src/cli/auth-broker-cli.ts @@ -0,0 +1,514 @@ +/** + * `omp auth-broker` command handlers. + * + * Sub-verbs: + * - `serve [--bind=…]` — boots the broker against the local SQLite store. + * - `token` / `token --regenerate` — manages the bearer token file. + * - `login [--via=user@host]` — logs into a provider locally, or + * via SSH tunnel into a remote broker host. + * - `import ` — imports CLIProxyAPI-style JSON credentials into + * the local SQLite store (typical use: `import ~/.cliproxy/auth`). + * - `status` — health-pings the configured remote broker. + */ +import * as crypto from "node:crypto"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { + AuthBrokerClient, + AuthStorage, + type CredentialDisabledEvent, + DEFAULT_AUTH_BROKER_BIND, + getOAuthProviders, + type OAuthCredential, + type OAuthProvider, + SqliteAuthCredentialStore, + startAuthBroker, +} from "@oh-my-pi/pi-ai"; +import { $which, APP_NAME, getAgentDbPath, getConfigRootDir, isEnoent, logger, VERSION } from "@oh-my-pi/pi-utils"; +import { $ } from "bun"; +import chalk from "chalk"; +import { resolveAuthBrokerConfig } from "../session/auth-broker-config"; + +export type AuthBrokerAction = "serve" | "token" | "login" | "logout" | "status" | "import"; + +export interface AuthBrokerCommandArgs { + action: AuthBrokerAction; + flags: { + json?: boolean; + bind?: string; + regenerate?: boolean; + via?: string; + provider?: string; + dryRun?: boolean; + /** `login`/`logout`: provider id. `import`: filesystem path. */ + source?: string; + /** `import`: keep credentials whose JSON had `disabled: true`. */ + includeDisabled?: boolean; + }; +} + +const ACTIONS: readonly AuthBrokerAction[] = ["serve", "token", "login", "logout", "import", "status"]; + +/** Callback ports baked from the per-provider OAuth flow modules. */ +const CALLBACK_PORTS: Record = { + anthropic: 54545, + "openai-codex": 1455, + "google-gemini-cli": 8085, + "google-antigravity": 51121, + "gitlab-duo": 8080, +}; + +function getTokenFilePath(): string { + return path.join(getConfigRootDir(), "auth-broker.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 Bun.write(file, token); + try { + await fs.chmod(file, 0o600); + } catch { + // Best-effort (e.g. Windows). + } +} + +function generateToken(): string { + return crypto.randomBytes(32).toString("base64url"); +} + +async function ensureToken(): Promise { + const existing = await readToken(); + if (existing) return existing; + const token = generateToken(); + await writeToken(token); + return token; +} + +async function runServe(flags: AuthBrokerCommandArgs["flags"]): Promise { + const bind = flags.bind ?? DEFAULT_AUTH_BROKER_BIND; + const token = await ensureToken(); + const dbPath = getAgentDbPath(); + const store = await SqliteAuthCredentialStore.open(dbPath); + const storage = new AuthStorage(store); + await storage.reload(); + const handle = startAuthBroker({ + storage, + bind, + bearerTokens: [token], + version: VERSION, + }); + process.stdout.write(`auth-broker listening on ${handle.url}\n`); + process.stdout.write(`bearer token: ${getTokenFilePath()} (chmod 0600)\n`); + + const credentialDisabledUnsub = storage.onCredentialDisabled((event: CredentialDisabledEvent) => { + logger.warn("auth-broker credential disabled", { ...event }); + }); + + const shutdown = async (signal: NodeJS.Signals): Promise => { + process.stdout.write(`\nReceived ${signal}, shutting down...\n`); + credentialDisabledUnsub(); + await handle.close(); + storage.close(); + process.exit(0); + }; + process.once("SIGINT", () => void shutdown("SIGINT")); + process.once("SIGTERM", () => void shutdown("SIGTERM")); + + // Block forever; lifecycle is signal-driven. + await new Promise(() => {}); +} + +async function runToken(flags: AuthBrokerCommandArgs["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 runLogin(flags: AuthBrokerCommandArgs["flags"]): Promise { + const providerArg = flags.provider; + if (!providerArg) { + throw new Error("Usage: omp auth-broker login [--via=user@host]"); + } + const oauthProviders = new Set(getOAuthProviders().map(p => p.id)); + if (!oauthProviders.has(providerArg)) { + throw new Error(`Unknown OAuth provider '${providerArg}'. Known: ${[...oauthProviders].sort().join(", ")}`); + } + if (flags.via) { + await runRemoteLogin(providerArg, flags.via, flags.dryRun ?? false); + return; + } + await runLocalLogin(providerArg as OAuthProvider); +} + +async function runLocalLogin(provider: OAuthProvider): Promise { + // Spawn the pi-ai CLI in-process — it handles the per-provider OAuth dance + // and persists into the same SQLite store the broker uses. + const piAiCli = Bun.fileURLToPath(import.meta.resolve("@oh-my-pi/pi-ai/cli")); + const proc = Bun.spawn({ + cmd: [process.execPath, piAiCli, "login", provider], + stdin: "inherit", + stdout: "inherit", + stderr: "inherit", + }); + const exitCode = await proc.exited; + if (exitCode !== 0) { + throw new Error(`pi-ai login exited with code ${exitCode}`); + } +} + +async function runRemoteLogin(provider: string, via: string, dryRun: boolean): Promise { + const port = CALLBACK_PORTS[provider]; + if (port === undefined) { + throw new Error( + `No known OAuth callback port for '${provider}'. Use device-code flow on the broker host directly.`, + ); + } + const sshArgs = [ + "-L", + `${port}:127.0.0.1:${port}`, + "-o", + "ExitOnForwardFailure=yes", + via, + `${APP_NAME} auth-broker login ${provider}`, + ]; + if (dryRun) { + process.stdout.write(`ssh ${sshArgs.map(a => (a.includes(" ") ? `'${a}'` : a)).join(" ")}\n`); + return; + } + const sshBin = $which("ssh"); + if (!sshBin) { + throw new Error("ssh binary not found in PATH"); + } + const proc = Bun.spawn({ + cmd: [sshBin, ...sshArgs], + stdin: "inherit", + stdout: "inherit", + stderr: "inherit", + }); + const exitCode = await proc.exited; + if (exitCode !== 0) { + throw new Error(`ssh exited with code ${exitCode}`); + } +} + +async function runLogout(flags: AuthBrokerCommandArgs["flags"]): Promise { + const providerArg = flags.provider; + if (!providerArg) { + throw new Error("Usage: omp auth-broker logout "); + } + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + store.deleteAuthCredentialsForProvider(providerArg, "logged out by user"); + process.stdout.write(`Logged out of ${providerArg}\n`); + } finally { + store.close(); + } +} + +// ─── CLIProxyAPI import ───────────────────────────────────────────────── + +/** + * Maps the `type` field of a CLIProxyAPI credential JSON to the omp provider id. + * The filename also encodes the type (e.g. `claude-foo@bar.json`), but the + * in-file `type` is authoritative — we only fall back to filename if absent. + */ +const CLIPROXY_TYPE_TO_PROVIDER: Record = { + claude: "anthropic", + codex: "openai-codex", + gemini: "google-gemini-cli", + antigravity: "google-antigravity", + "gemini-cli": "google-gemini-cli", +}; + +interface CliProxyCredentialJson { + type?: string; + access_token?: string; + refresh_token?: string; + id_token?: string; + expired?: string; + last_refresh?: string; + email?: string; + account_id?: string; + disabled?: boolean; +} + +interface ImportPlanEntry { + sourceFile: string; + provider: string; + email: string | null; + accountId: string | null; + expiresAt: number; + disabled: boolean; + credential: OAuthCredential; +} + +function resolveCliProxyProvider(json: CliProxyCredentialJson, filename: string, overrideId?: string): string | null { + if (overrideId && overrideId.length > 0) return overrideId; + const typeField = json.type?.trim().toLowerCase(); + if (typeField && CLIPROXY_TYPE_TO_PROVIDER[typeField]) return CLIPROXY_TYPE_TO_PROVIDER[typeField]; + // Fall back to filename prefix: `-.json` + const base = path.basename(filename, ".json").toLowerCase(); + for (const prefix in CLIPROXY_TYPE_TO_PROVIDER) { + const providerId = CLIPROXY_TYPE_TO_PROVIDER[prefix]; + if (base.startsWith(`${prefix}-`) || base === prefix) return providerId; + } + return null; +} + +function parseCliProxyExpiry(raw: string | undefined): number | null { + if (!raw) return null; + // CLIProxyAPI writes RFC3339-ish dates. `Date.parse` handles both `Z` and offsets. + const ms = Date.parse(raw); + if (!Number.isFinite(ms)) return null; + return ms; +} + +async function collectImportSources(target: string): Promise { + const stat = await fs.stat(target); + if (stat.isFile()) return [target]; + if (!stat.isDirectory()) { + throw new Error(`Import source is neither file nor directory: ${target}`); + } + const entries = await fs.readdir(target, { withFileTypes: true }); + const files: string[] = []; + for (const entry of entries) { + if (!entry.isFile()) continue; + if (!entry.name.endsWith(".json")) continue; + files.push(path.join(target, entry.name)); + } + files.sort(); + return files; +} + +async function loadImportPlan( + target: string, + overrideProvider: string | undefined, + includeDisabled: boolean, +): Promise<{ entries: ImportPlanEntry[]; skipped: Array<{ file: string; reason: string }> }> { + const files = await collectImportSources(target); + const entries: ImportPlanEntry[] = []; + const skipped: Array<{ file: string; reason: string }> = []; + for (const file of files) { + let json: CliProxyCredentialJson; + try { + json = (await Bun.file(file).json()) as CliProxyCredentialJson; + } catch (err) { + skipped.push({ file, reason: `unreadable JSON: ${String(err)}` }); + continue; + } + if (json.disabled === true && !includeDisabled) { + skipped.push({ file, reason: "credential marked disabled (use --include-disabled to import anyway)" }); + continue; + } + const provider = resolveCliProxyProvider(json, file, overrideProvider); + if (!provider) { + skipped.push({ + file, + reason: `cannot determine omp provider from type=${json.type ?? "?"} (pass --provider to override)`, + }); + continue; + } + if (!json.access_token || !json.refresh_token) { + skipped.push({ file, reason: "missing access_token or refresh_token" }); + continue; + } + const expiresAt = parseCliProxyExpiry(json.expired); + if (expiresAt === null) { + skipped.push({ file, reason: `cannot parse expired=${json.expired ?? "?"}` }); + continue; + } + const email = typeof json.email === "string" && json.email.length > 0 ? json.email : null; + const accountId = typeof json.account_id === "string" && json.account_id.length > 0 ? json.account_id : null; + const credential: OAuthCredential = { + type: "oauth", + access: json.access_token, + refresh: json.refresh_token, + expires: expiresAt, + ...(email !== null ? { email } : {}), + ...(accountId !== null ? { accountId } : {}), + }; + entries.push({ + sourceFile: file, + provider, + email, + accountId, + expiresAt, + disabled: json.disabled === true, + credential, + }); + } + return { entries, skipped }; +} + +function describeImportEntry(entry: ImportPlanEntry): string { + const ident = entry.email ?? entry.accountId ?? "(no identity)"; + const stale = entry.expiresAt < Date.now() ? " [expired]" : ""; + const disabled = entry.disabled ? " [disabled]" : ""; + return `${entry.provider}: ${ident}${stale}${disabled} from ${entry.sourceFile}`; +} + +async function runImport(flags: AuthBrokerCommandArgs["flags"]): Promise { + const target = flags.source; + if (!target) { + throw new Error("Usage: omp auth-broker import [--provider=] [--include-disabled] [--dry-run]"); + } + const resolvedTarget = path.resolve(target.startsWith("~") ? target.replace(/^~/, os.homedir()) : target); + const { entries, skipped } = await loadImportPlan(resolvedTarget, flags.provider, flags.includeDisabled === true); + + if (flags.json) { + process.stdout.write( + `${JSON.stringify({ + dryRun: flags.dryRun === true, + imported: flags.dryRun + ? [] + : entries.map(e => ({ provider: e.provider, email: e.email, file: e.sourceFile })), + plan: entries.map(e => ({ + provider: e.provider, + email: e.email, + accountId: e.accountId, + expiresAt: e.expiresAt, + disabled: e.disabled, + file: e.sourceFile, + })), + skipped, + })}\n`, + ); + } + + if (!flags.json) { + for (const skip of skipped) { + process.stdout.write(`${chalk.yellow("skip")} ${skip.file}: ${skip.reason}\n`); + } + } + + if (entries.length === 0) { + if (!flags.json) process.stdout.write(`No importable credentials in ${resolvedTarget}.\n`); + return; + } + + if (flags.dryRun === true) { + if (!flags.json) { + process.stdout.write(`Dry run — would import ${entries.length} credential(s):\n`); + for (const entry of entries) process.stdout.write(` ${describeImportEntry(entry)}\n`); + } + return; + } + + const brokerConfig = await resolveAuthBrokerConfig(); + if (brokerConfig) { + const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token }); + for (const entry of entries) { + try { + await client.uploadCredential(entry.provider, entry.credential); + if (!flags.json) { + process.stdout.write(`${chalk.green("uploaded")} ${describeImportEntry(entry)} → ${brokerConfig.url}\n`); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + if (flags.json) { + process.stdout.write(`${JSON.stringify({ error: message, file: entry.sourceFile })}\n`); + } else { + process.stdout.write(`${chalk.red("failed")} ${describeImportEntry(entry)}: ${message}\n`); + } + process.exitCode = 1; + } + } + return; + } + + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + for (const entry of entries) { + store.upsertAuthCredentialForProvider(entry.provider, entry.credential); + if (!flags.json) process.stdout.write(`${chalk.green("imported")} ${describeImportEntry(entry)}\n`); + } + } finally { + store.close(); + } +} + +async function runStatus(flags: AuthBrokerCommandArgs["flags"]): Promise { + const cfg = await resolveAuthBrokerConfig(); + if (!cfg) { + const message = "No auth-broker configured (set OMP_AUTH_BROKER_URL to enable)."; + if (flags.json) process.stdout.write(`${JSON.stringify({ ok: false, reason: "not_configured" })}\n`); + else process.stdout.write(`${chalk.yellow(message)}\n`); + return; + } + const client = new AuthBrokerClient({ url: cfg.url, token: cfg.token }); + try { + const health = await client.healthz(); + if (flags.json) { + process.stdout.write(`${JSON.stringify({ url: cfg.url, ...health })}\n`); + } else { + process.stdout.write(`${chalk.green("OK")} ${cfg.url} (version=${health.version ?? "unknown"})\n`); + } + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + if (flags.json) { + process.stdout.write(`${JSON.stringify({ ok: false, url: cfg.url, error: message })}\n`); + } else { + process.stdout.write(`${chalk.red("FAILED")} ${cfg.url}: ${message}\n`); + } + process.exitCode = 1; + } +} + +export async function runAuthBrokerCommand(cmd: AuthBrokerCommandArgs): Promise { + switch (cmd.action) { + case "serve": + await runServe(cmd.flags); + return; + case "token": + await runToken(cmd.flags); + return; + case "login": + await runLogin(cmd.flags); + return; + case "logout": + await runLogout(cmd.flags); + return; + case "import": + await runImport(cmd.flags); + return; + case "status": + await runStatus(cmd.flags); + return; + default: { + // Exhaustive check. + const _exhaustive: never = cmd.action; + throw new Error(`Unknown auth-broker action: ${String(_exhaustive)}`); + } + } +} + +export { ACTIONS as AUTH_BROKER_ACTIONS }; + +// Touch `$` so Bun's tree-shaker keeps the shell helper imported (used by future verbs). +void $; diff --git a/packages/coding-agent/src/commands/auth-broker.ts b/packages/coding-agent/src/commands/auth-broker.ts new file mode 100644 index 000000000..ddb250836 --- /dev/null +++ b/packages/coding-agent/src/commands/auth-broker.ts @@ -0,0 +1,82 @@ +/** + * `omp auth-broker` — manage the omp credential vault. + */ +import { Args, Command, Flags, renderCommandHelp } from "@oh-my-pi/pi-utils/cli"; +import { + AUTH_BROKER_ACTIONS, + type AuthBrokerAction, + type AuthBrokerCommandArgs, + runAuthBrokerCommand, +} from "../cli/auth-broker-cli"; +import { initTheme } from "../modes/theme/theme"; + +export default class AuthBroker extends Command { + static description = "Manage the omp auth-broker (credential vault)"; + + static args = { + action: Args.string({ + description: "Sub-command", + required: false, + options: [...AUTH_BROKER_ACTIONS], + }), + // Second positional: provider id (login/logout) or filesystem path (import). + source: Args.string({ + description: "OAuth provider id (login/logout) or path (import)", + required: false, + }), + }; + + static flags = { + json: Flags.boolean({ description: "Output JSON" }), + bind: Flags.string({ description: "Bind address for `serve` (host:port)", char: "b" }), + regenerate: Flags.boolean({ description: "Regenerate the bearer token" }), + via: Flags.string({ + description: "SSH user@host for remote login (login --via=user@host)", + }), + provider: Flags.string({ + description: "Override provider id for `import` (e.g. when JSON `type` is unrecognized)", + }), + "include-disabled": Flags.boolean({ + description: "Import credentials whose JSON has `disabled: true` (import)", + }), + "dry-run": Flags.boolean({ description: "Print actions without executing (import / login --via)" }), + }; + + static examples = [ + "# Boot the broker against the local SQLite store\n omp auth-broker serve", + "# Boot on a non-default port\n omp auth-broker serve --bind=127.0.0.1:9000", + "# Print the bearer token\n omp auth-broker token", + "# Rotate the bearer token\n omp auth-broker token --regenerate", + "# Local login (run on the broker host)\n omp auth-broker login anthropic", + "# Remote login over SSH tunnel\n omp auth-broker login anthropic --via=user@broker", + "# Import a CLIProxyAPI auth dump\n omp auth-broker import ~/.cliproxy/auth", + "# Import a single CLIProxyAPI JSON, overriding the provider mapping\n omp auth-broker import ~/.cliproxy/auth/claude-foo.json --provider anthropic", + "# Health-check the configured remote broker\n omp auth-broker status", + ]; + + async run(): Promise { + const { args, flags } = await this.parse(AuthBroker); + if (!args.action) { + renderCommandHelp("omp", "auth-broker", AuthBroker); + return; + } + const action = args.action as AuthBrokerAction; + const cmd: AuthBrokerCommandArgs = { + action, + flags: { + json: flags.json, + bind: flags.bind, + regenerate: flags.regenerate, + via: flags.via, + // `login`/`logout` reuse the legacy `provider` slot; `import` keeps `source` separate + // so `provider` flag (used as an override) is unambiguous. + provider: action === "import" ? flags.provider : (args.source ?? flags.provider), + source: args.source, + includeDisabled: flags["include-disabled"], + dryRun: flags["dry-run"], + }, + }; + await initTheme(); + await runAuthBrokerCommand(cmd); + } +} diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index 289a1756f..cf9e877fc 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -234,6 +234,13 @@ export const SETTINGS_SCHEMA = { // ──────────────────────────────────────────────────────────────────────── lastChangelogVersion: { type: "string", default: undefined }, + // Auth broker — credentials proxied through a remote `omp auth-broker serve` + // host. Hidden from the UI; populate via env vars or hand-edited config.yml. + // Env (`OMP_AUTH_BROKER_URL` / `OMP_AUTH_BROKER_TOKEN`) takes precedence so + // per-machine overrides remain trivial. + "auth.broker.url": { type: "string", default: undefined }, + "auth.broker.token": { type: "string", default: undefined }, + autoResume: { type: "boolean", default: false, diff --git a/packages/coding-agent/src/modes/controllers/command-controller.ts b/packages/coding-agent/src/modes/controllers/command-controller.ts index 18ff6acca..9f426b675 100644 --- a/packages/coding-agent/src/modes/controllers/command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/command-controller.ts @@ -374,10 +374,15 @@ export class CommandController { const openaiWebsocketSetting = this.ctx.settings.get("providers.openaiWebsockets") ?? "auto"; const preferOpenAICodexWebsockets = openaiWebsocketSetting === "on" ? true : openaiWebsocketSetting === "off" ? false : undefined; + const credentialSource = this.ctx.session.modelRegistry.authStorage.describeCredentialSource( + model.provider, + stats.sessionId, + ); const providerDetails = getProviderDetails({ model, sessionId: stats.sessionId, authMode, + credentialSource, preferWebsockets: preferOpenAICodexWebsockets, providerSessionState: this.ctx.session.providerSessionState, }); diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index cbbaad3ab..e5ca3f646 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -93,7 +93,13 @@ import { SecretObfuscator, } from "./secrets"; import { AgentSession } from "./session/agent-session"; -import { AuthStorage } from "./session/auth-storage"; +import { resolveAuthBrokerConfig } from "./session/auth-broker-config"; +import { + AuthBrokerClient, + AuthStorage, + REMOTE_REFRESH_SENTINEL, + RemoteAuthCredentialStore, +} from "./session/auth-storage"; import { convertToLlm } from "./session/messages"; import { SessionManager } from "./session/session-manager"; import { closeAllConnections } from "./ssh/connection-manager"; @@ -317,13 +323,52 @@ function getDefaultAgentDir(): string { // Discovery Functions /** - * Create an AuthStorage instance with fallback support. - * Reads from primary path first, then falls back to legacy paths (.pi, .claude). + * Create an AuthStorage instance. + * + * Default: local SQLite store at `/agent.db`. + * + * Broker mode: when `OMP_AUTH_BROKER_URL` is set, credentials are pulled from + * a remote auth-broker over the wire. Refresh tokens never leave the broker; + * the client receives access tokens with `refresh = "__remote__"` and calls + * back into the broker through the {@link AuthStorageOptions.refreshOAuthCredential} + * override to re-mint access tokens when needed. */ export async function discoverAuthStorage(agentDir: string = getDefaultAgentDir()): Promise { + const brokerConfig = await resolveAuthBrokerConfig(); + if (brokerConfig) { + const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token }); + const initialSnapshot = await client.fetchSnapshot(); + const store = new RemoteAuthCredentialStore({ client, initialSnapshot }); + const storage = new AuthStorage(store, { + configValueResolver: resolveConfigValue, + sourceLabel: `broker ${brokerConfig.url}`, + refreshOAuthCredential: async (_provider, credentialId, _credential) => { + const { entry } = await client.refreshCredential(credentialId); + if (entry.credential.type !== "oauth") { + throw new Error(`Broker returned non-OAuth credential for id=${credentialId}`); + } + const refreshed = entry.credential; + return { + access: refreshed.access, + // Sentinel — AuthStorage stores it back into the in-memory snapshot, + // but a refresh through the broker is the only legal way to mint tokens. + refresh: REMOTE_REFRESH_SENTINEL, + expires: refreshed.expires, + accountId: refreshed.accountId, + email: refreshed.email, + projectId: refreshed.projectId, + enterpriseUrl: refreshed.enterpriseUrl, + }; + }, + }); + await storage.reload(); + return storage; + } const dbPath = getAgentDbPath(agentDir); - - const storage = await AuthStorage.create(dbPath, { configValueResolver: resolveConfigValue }); + const storage = await AuthStorage.create(dbPath, { + configValueResolver: resolveConfigValue, + sourceLabel: `local ${dbPath}`, + }); await storage.reload(); return storage; } diff --git a/packages/coding-agent/src/session/auth-broker-config.ts b/packages/coding-agent/src/session/auth-broker-config.ts new file mode 100644 index 000000000..33d543050 --- /dev/null +++ b/packages/coding-agent/src/session/auth-broker-config.ts @@ -0,0 +1,102 @@ +/** + * Resolve auth-broker connection configuration for the local omp client. + * + * Precedence (highest first): + * 1. `OMP_AUTH_BROKER_URL` / `OMP_AUTH_BROKER_TOKEN` env vars. + * 2. `auth.broker.url` / `auth.broker.token` in `~/.omp/agent/config.yml` + * (hidden from the settings UI; `!command` resolution supported). + * 3. Token file `~/.omp/auth-broker.token` (paired with URL from env or config). + * + * Returns null when no broker URL is configured — caller falls back to the + * local SQLite store. + * + * Reads config.yml directly (instead of going through `Settings.init`) because + * `discoverAuthStorage` runs before the settings singleton is initialized in + * `runRootCommand`, and we want hand-edited config entries to be honoured at + * boot without forcing a startup reorder. + */ +import * as path from "node:path"; +import { getAgentDir, getConfigRootDir, isEnoent, logger } from "@oh-my-pi/pi-utils"; +import { YAML } from "bun"; +import { resolveConfigValue } from "../config/resolve-config-value"; + +export interface AuthBrokerClientConfig { + url: string; + token: string; +} + +/** Path to the local bearer token file. Created on the broker host by `omp auth-broker token`. */ +export function getAuthBrokerTokenFilePath(): string { + return path.join(getConfigRootDir(), "auth-broker.token"); +} + +async function readTokenFile(): Promise { + try { + const raw = await Bun.file(getAuthBrokerTokenFilePath()).text(); + const trimmed = raw.trim(); + return trimmed.length > 0 ? trimmed : null; + } catch (err) { + if (isEnoent(err)) return null; + logger.warn("auth-broker token file unreadable", { error: String(err) }); + return null; + } +} + +interface ConfigSnapshot { + url?: string; + token?: string; +} + +async function readConfigYaml(): Promise { + const configPath = path.join(getAgentDir(), "config.yml"); + try { + const raw = await Bun.file(configPath).text(); + const parsed = YAML.parse(raw); + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return {}; + const record = parsed as Record; + const url = typeof record["auth.broker.url"] === "string" ? (record["auth.broker.url"] as string) : undefined; + const token = + typeof record["auth.broker.token"] === "string" ? (record["auth.broker.token"] as string) : undefined; + return { url, token }; + } catch (err) { + if (isEnoent(err)) return {}; + logger.warn("auth-broker config.yml unreadable", { error: String(err) }); + return {}; + } +} + +/** + * Read broker configuration. Returns null when the URL is missing + * (broker disabled — local store is used). Throws when URL is set but no + * token is available — the caller cannot fall back silently because the + * user explicitly asked to use the broker. + */ +export async function resolveAuthBrokerConfig(): Promise { + const envUrl = process.env.OMP_AUTH_BROKER_URL; + const envToken = process.env.OMP_AUTH_BROKER_TOKEN; + + let url = envUrl && envUrl.length > 0 ? envUrl : undefined; + let configToken: string | undefined; + if (!url || !envToken) { + const fromConfig = await readConfigYaml(); + if (!url && fromConfig.url) { + const resolved = await resolveConfigValue(fromConfig.url); + if (resolved && resolved.length > 0) url = resolved; + } + if (fromConfig.token) { + const resolved = await resolveConfigValue(fromConfig.token); + if (resolved && resolved.length > 0) configToken = resolved; + } + } + if (!url) return null; + + const token = + (envToken && envToken.length > 0 ? envToken : undefined) ?? configToken ?? (await readTokenFile()) ?? undefined; + if (!token) { + throw new Error( + `OMP_AUTH_BROKER_URL is set (${url}) but no bearer token is available. ` + + `Set OMP_AUTH_BROKER_TOKEN, the \`auth.broker.token\` config entry, or place one at ${getAuthBrokerTokenFilePath()}.`, + ); + } + return { url, token }; +} diff --git a/packages/coding-agent/src/session/auth-storage.ts b/packages/coding-agent/src/session/auth-storage.ts index 49d670eae..33f0d1607 100644 --- a/packages/coding-agent/src/session/auth-storage.ts +++ b/packages/coding-agent/src/session/auth-storage.ts @@ -14,4 +14,10 @@ export type { SerializedAuthStorage, StoredAuthCredential, } from "@oh-my-pi/pi-ai"; -export { AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai"; +export { + AuthBrokerClient, + AuthStorage, + REMOTE_REFRESH_SENTINEL, + RemoteAuthCredentialStore, + SqliteAuthCredentialStore, +} from "@oh-my-pi/pi-ai"; diff --git a/packages/coding-agent/test/auth-broker-import.test.ts b/packages/coding-agent/test/auth-broker-import.test.ts new file mode 100644 index 000000000..877210325 --- /dev/null +++ b/packages/coding-agent/test/auth-broker-import.test.ts @@ -0,0 +1,303 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { type AuthBrokerServerHandle, AuthStorage, SqliteAuthCredentialStore, startAuthBroker } from "@oh-my-pi/pi-ai"; +import { getAgentDbPath, setAgentDir } from "@oh-my-pi/pi-utils"; +import { runAuthBrokerCommand } from "../src/cli/auth-broker-cli"; + +const ORIGINAL_STDOUT_WRITE = process.stdout.write.bind(process.stdout); + +function silenceStdout(): () => string { + let captured = ""; + process.stdout.write = ((chunk: string | Uint8Array): boolean => { + captured += typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); + return true; + }) as typeof process.stdout.write; + return () => captured; +} + +describe("auth-broker import (CLIProxyAPI)", () => { + let agentDir = ""; + let cliproxyDir = ""; + let originalAgentDir: string | undefined; + + beforeEach(async () => { + originalAgentDir = process.env.OMP_AGENT_DIR; + agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-import-agent-")); + cliproxyDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-import-cliproxy-")); + setAgentDir(agentDir); + }); + + afterEach(async () => { + process.stdout.write = ORIGINAL_STDOUT_WRITE; + if (originalAgentDir === undefined) delete process.env.OMP_AGENT_DIR; + else process.env.OMP_AGENT_DIR = originalAgentDir; + await fs.rm(agentDir, { recursive: true, force: true }); + await fs.rm(cliproxyDir, { recursive: true, force: true }); + }); + + async function writeCliProxyJson(name: string, body: Record): Promise { + const file = path.join(cliproxyDir, name); + await Bun.write(file, JSON.stringify(body)); + return file; + } + + test("imports a directory of CLIProxyAPI JSONs and maps types to omp providers", async () => { + await writeCliProxyJson("claude-sample.json", { + type: "claude", + access_token: "claude-access-1", + refresh_token: "claude-refresh-1", + expired: "2099-12-31T23:59:59Z", + email: "claude-user@example.com", + id_token: "ignored", + last_refresh: "2025-01-01T00:00:00Z", + }); + await writeCliProxyJson("codex-sample.json", { + type: "codex", + access_token: "codex-access-1", + refresh_token: "codex-refresh-1", + expired: "2099-12-31T23:59:59Z", + email: "codex-user@example.com", + account_id: "acct-codex-1", + websockets: true, + }); + await writeCliProxyJson("disabled.json", { + type: "claude", + access_token: "x", + refresh_token: "y", + expired: "2099-12-31T23:59:59Z", + email: "disabled@example.com", + disabled: true, + }); + + const restore = silenceStdout(); + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir, json: false }, + }); + restore(); + + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + const claude = store.listAuthCredentials("anthropic"); + expect(claude).toHaveLength(1); + expect(claude[0].credential.type).toBe("oauth"); + if (claude[0].credential.type === "oauth") { + expect(claude[0].credential.access).toBe("claude-access-1"); + expect(claude[0].credential.refresh).toBe("claude-refresh-1"); + expect(claude[0].credential.email).toBe("claude-user@example.com"); + expect(claude[0].credential.expires).toBe(Date.parse("2099-12-31T23:59:59Z")); + } + + const codex = store.listAuthCredentials("openai-codex"); + expect(codex).toHaveLength(1); + if (codex[0].credential.type === "oauth") { + expect(codex[0].credential.access).toBe("codex-access-1"); + expect(codex[0].credential.accountId).toBe("acct-codex-1"); + } + + // disabled.json was skipped by default + const disabled = store + .listAuthCredentials("anthropic") + .find(r => r.credential.type === "oauth" && r.credential.email === "disabled@example.com"); + expect(disabled).toBeUndefined(); + } finally { + store.close(); + } + }); + + test("dry-run does not write any credentials", async () => { + await writeCliProxyJson("claude.json", { + type: "claude", + access_token: "a", + refresh_token: "b", + expired: "2099-12-31T23:59:59Z", + email: "dryrun@example.com", + }); + + const restore = silenceStdout(); + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir, dryRun: true, json: true }, + }); + const output = restore(); + + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + expect(store.listAuthCredentials()).toHaveLength(0); + } finally { + store.close(); + } + const parsed = JSON.parse(output.trim().split("\n").pop() ?? "{}"); + expect(parsed.dryRun).toBe(true); + expect(parsed.plan).toHaveLength(1); + expect(parsed.plan[0].provider).toBe("anthropic"); + }); + + test("--provider override forces a provider id when the JSON type is unrecognized", async () => { + await writeCliProxyJson("weird.json", { + type: "some-future-type", + access_token: "z", + refresh_token: "w", + expired: "2099-12-31T23:59:59Z", + email: "future@example.com", + }); + + const restore = silenceStdout(); + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir, provider: "anthropic" }, + }); + restore(); + + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + const rows = store.listAuthCredentials("anthropic"); + expect(rows).toHaveLength(1); + } finally { + store.close(); + } + }); + + test("--include-disabled imports rows marked disabled", async () => { + await writeCliProxyJson("disabled.json", { + type: "claude", + access_token: "d", + refresh_token: "e", + expired: "2099-12-31T23:59:59Z", + email: "disabled-import@example.com", + disabled: true, + }); + + const restore = silenceStdout(); + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir, includeDisabled: true }, + }); + restore(); + + const store = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + expect(store.listAuthCredentials("anthropic")).toHaveLength(1); + } finally { + store.close(); + } + }); +}); + +describe("auth-broker import (broker-routed)", () => { + let agentDir = ""; + let brokerAgentDir = ""; + let cliproxyDir = ""; + let brokerStore: SqliteAuthCredentialStore | undefined; + let brokerStorage: AuthStorage | undefined; + let handle: AuthBrokerServerHandle | undefined; + const token = "broker-import-bearer"; + const savedEnv: Record = {}; + + beforeEach(async () => { + savedEnv.OMP_AUTH_BROKER_URL = process.env.OMP_AUTH_BROKER_URL; + savedEnv.OMP_AUTH_BROKER_TOKEN = process.env.OMP_AUTH_BROKER_TOKEN; + agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-import-client-")); + brokerAgentDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-import-broker-")); + cliproxyDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-import-cliproxy-broker-")); + setAgentDir(agentDir); + + brokerStore = await SqliteAuthCredentialStore.open(path.join(brokerAgentDir, "agent.db")); + brokerStorage = new AuthStorage(brokerStore); + await brokerStorage.reload(); + handle = startAuthBroker({ + storage: brokerStorage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + process.env.OMP_AUTH_BROKER_URL = handle.url; + process.env.OMP_AUTH_BROKER_TOKEN = token; + }); + + afterEach(async () => { + await handle?.close(); + brokerStorage?.close(); + brokerStore?.close(); + await fs.rm(agentDir, { recursive: true, force: true }); + await fs.rm(brokerAgentDir, { recursive: true, force: true }); + await fs.rm(cliproxyDir, { recursive: true, force: true }); + for (const key of ["OMP_AUTH_BROKER_URL", "OMP_AUTH_BROKER_TOKEN"] as const) { + if (savedEnv[key] === undefined) delete process.env[key]; + else process.env[key] = savedEnv[key]; + } + }); + + test("uploads CLIProxyAPI JSONs to the broker when configured, not the local store", async () => { + await Bun.write( + path.join(cliproxyDir, "claude-foo@bar.json"), + JSON.stringify({ + type: "claude", + access_token: "broker-access", + refresh_token: "broker-refresh-real", + expired: "2099-12-31T23:59:59Z", + email: "foo@bar.com", + }), + ); + + const ORIGINAL_STDOUT = process.stdout.write.bind(process.stdout); + let captured = ""; + process.stdout.write = ((chunk: string | Uint8Array): boolean => { + captured += typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); + return true; + }) as typeof process.stdout.write; + try { + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir }, + }); + } finally { + process.stdout.write = ORIGINAL_STDOUT; + } + + // The broker received it (and persisted the real refresh token). + const persisted = brokerStore!.getOAuth("anthropic"); + expect(persisted?.access).toBe("broker-access"); + expect(persisted?.refresh).toBe("broker-refresh-real"); + expect(persisted?.email).toBe("foo@bar.com"); + + // The local client SQLite store was NOT touched. + const localStore = await SqliteAuthCredentialStore.open(getAgentDbPath()); + try { + expect(localStore.listAuthCredentials()).toHaveLength(0); + } finally { + localStore.close(); + } + + expect(captured).toContain("uploaded"); + expect(captured).toContain(handle!.url); + }); + + test("dry-run does not upload even when broker is configured", async () => { + await Bun.write( + path.join(cliproxyDir, "claude-dry.json"), + JSON.stringify({ + type: "claude", + access_token: "a", + refresh_token: "b", + expired: "2099-12-31T23:59:59Z", + email: "dry@example.com", + }), + ); + + const ORIGINAL_STDOUT = process.stdout.write.bind(process.stdout); + process.stdout.write = (() => true) as typeof process.stdout.write; + try { + await runAuthBrokerCommand({ + action: "import", + flags: { source: cliproxyDir, dryRun: true }, + }); + } finally { + process.stdout.write = ORIGINAL_STDOUT; + } + + expect(brokerStore!.listAuthCredentials()).toHaveLength(0); + }); +}); diff --git a/packages/utils/src/logger.ts b/packages/utils/src/logger.ts index 7b7270962..10a81f963 100644 --- a/packages/utils/src/logger.ts +++ b/packages/utils/src/logger.ts @@ -84,6 +84,19 @@ export function warn(message: string, context?: Record): void { } } +/** + * Log an informational message. + * @param message - The message to log. + * @param context - The context to log. + */ +export function info(message: string, context?: Record): void { + try { + winstonLogger.info(message, context); + } catch { + // Silently ignore logging failures + } +} + /** * Log a debug message. * @param message - The message to log.