@oh-my-pi/pi-ai 6.7.67 → 6.8.0

package/src/storage.ts

@@ -0,0 +1,185 @@
+/**
+ * Simple auth storage for CLI using SQLite.
+ * Compatible with coding-agent's agent.db format.
+ */
+
+import { Database } from "bun:sqlite";
+import { existsSync, mkdirSync } from "node:fs";
+import { homedir } from "node:os";
+import { dirname, join } from "node:path";
+import type { OAuthCredentials } from "./utils/oauth/types";
+
+type AuthCredential = { type: "api_key"; key: string } | ({ type: "oauth" } & OAuthCredentials);
+
+type AuthRow = {
+	id: number;
+	provider: string;
+	credential_type: string;
+	data: string;
+	created_at: number;
+	updated_at: number;
+};
+
+/**
+ * Get the agent config directory (e.g., ~/.omp/agent/)
+ */
+function getAgentDir(): string {
+	const configDir = process.env.OMP_CODING_AGENT_DIR || join(homedir(), ".omp", "agent");
+	return configDir;
+}
+
+/**
+ * Get path to agent.db
+ */
+function getAgentDbPath(): string {
+	return join(getAgentDir(), "agent.db");
+}
+
+function serializeCredential(credential: AuthCredential): { credentialType: string; data: string } | null {
+	if (credential.type === "api_key") {
+		return {
+			credentialType: "api_key",
+			data: JSON.stringify({ key: credential.key }),
+		};
+	}
+	if (credential.type === "oauth") {
+		const { type: _type, ...rest } = credential;
+		return {
+			credentialType: "oauth",
+			data: JSON.stringify(rest),
+		};
+	}
+	return null;
+}
+
+function deserializeCredential(row: AuthRow): AuthCredential | null {
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(row.data);
+	} catch {
+		return null;
+	}
+	if (!parsed || typeof parsed !== "object") {
+		return null;
+	}
+
+	if (row.credential_type === "api_key") {
+		const data = parsed as Record<string, unknown>;
+		if (typeof data.key === "string") {
+			return { type: "api_key", key: data.key };
+		}
+	}
+
+	if (row.credential_type === "oauth") {
+		return { type: "oauth", ...(parsed as Record<string, unknown>) } as AuthCredential;
+	}
+
+	return null;
+}
+
+/**
+ * Simple storage class for CLI auth credentials.
+ */
+export class CliAuthStorage {
+	private db: Database;
+	private insertStmt: ReturnType<Database["prepare"]>;
+	private listByProviderStmt: ReturnType<Database["prepare"]>;
+	private listAllStmt: ReturnType<Database["prepare"]>;
+	private deleteByProviderStmt: ReturnType<Database["prepare"]>;
+
+	constructor(dbPath: string = getAgentDbPath()) {
+		// Ensure directory exists
+		const dir = dirname(dbPath);
+		if (!existsSync(dir)) {
+			mkdirSync(dir, { recursive: true });
+		}
+
+		this.db = new Database(dbPath);
+		this.initializeSchema();
+
+		this.insertStmt = this.db.prepare(
+			"INSERT INTO auth_credentials (provider, credential_type, data) VALUES (?, ?, ?) RETURNING id",
+		);
+		this.listByProviderStmt = this.db.prepare("SELECT * FROM auth_credentials WHERE provider = ?");
+		this.listAllStmt = this.db.prepare("SELECT * FROM auth_credentials");
+		this.deleteByProviderStmt = this.db.prepare("DELETE FROM auth_credentials WHERE provider = ?");
+	}
+
+	private initializeSchema(): void {
+		this.db.exec(`
+PRAGMA journal_mode=WAL;
+PRAGMA synchronous=NORMAL;
+PRAGMA busy_timeout=5000;
+
+CREATE TABLE IF NOT EXISTS auth_credentials (
+	id INTEGER PRIMARY KEY AUTOINCREMENT,
+	provider TEXT NOT NULL,
+	credential_type TEXT NOT NULL,
+	data TEXT NOT NULL,
+	created_at INTEGER NOT NULL DEFAULT (unixepoch()),
+	updated_at INTEGER NOT NULL DEFAULT (unixepoch())
+);
+CREATE INDEX IF NOT EXISTS idx_auth_provider ON auth_credentials(provider);
+		`);
+	}
+
+	/**
+	 * Save OAuth credentials for a provider (replaces existing).
+	 */
+	saveOAuth(provider: string, credentials: OAuthCredentials): void {
+		const credential: AuthCredential = { type: "oauth", ...credentials };
+		this.replaceForProvider(provider, credential);
+	}
+
+	/**
+	 * Get OAuth credentials for a provider.
+	 */
+	getOAuth(provider: string): OAuthCredentials | null {
+		const rows = this.listByProviderStmt.all(provider) as AuthRow[];
+		for (const row of rows) {
+			const credential = deserializeCredential(row);
+			if (credential && credential.type === "oauth") {
+				const { type: _type, ...oauth } = credential;
+				return oauth as OAuthCredentials;
+			}
+		}
+		return null;
+	}
+
+	/**
+	 * List all providers with credentials.
+	 */
+	listProviders(): string[] {
+		const rows = this.listAllStmt.all() as AuthRow[];
+		const providers = new Set<string>();
+		for (const row of rows) {
+			providers.add(row.provider);
+		}
+		return Array.from(providers);
+	}
+
+	/**
+	 * Delete all credentials for a provider.
+	 */
+	deleteProvider(provider: string): void {
+		this.deleteByProviderStmt.run(provider);
+	}
+
+	/**
+	 * Replace all credentials for a provider with a single credential.
+	 */
+	private replaceForProvider(provider: string, credential: AuthCredential): void {
+		const serialized = serializeCredential(credential);
+		if (!serialized) return;
+
+		const replace = this.db.transaction(() => {
+			this.deleteByProviderStmt.run(provider);
+			this.insertStmt.run(provider, serialized.credentialType, serialized.data);
+		});
+		replace();
+	}
+
+	close(): void {
+		this.db.close();
+	}
+}

package/src/utils/event-stream.ts

@@ -12,9 +12,9 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 		private isComplete: (event: T) => boolean,
 		private extractResult: (event: T) => R,
 	) {
-		this.finalResultPromise = new Promise((resolve) => {
-			this.resolveFinalResult = resolve;
-		});
+		const { promise, resolve } = Promise.withResolvers<R>();
+		this.finalResultPromise = promise;
+		this.resolveFinalResult = resolve;
 	}
 
 	push(event: T): void {

package/src/utils/oauth/anthropic.ts

@@ -2,109 +2,91 @@
  * Anthropic OAuth flow (Claude Pro/Max)
  */
 
+import { OAuthCallbackFlow } from "./callback-server";
 import { generatePKCE } from "./pkce";
-import type { OAuthCredentials } from "./types";
+import type { OAuthController, OAuthCredentials } from "./types";
 
 const decode = (s: string) => atob(s);
 const CLIENT_ID = decode("OWQxYzI1MGEtZTYxYi00NGQ5LTg4ZWQtNTk0NGQxOTYyZjVl");
 const AUTHORIZE_URL = "https://claude.ai/oauth/authorize";
 const TOKEN_URL = "https://console.anthropic.com/v1/oauth/token";
-const REDIRECT_URI = "https://console.anthropic.com/oauth/code/callback";
+const CALLBACK_PORT = 54545;
+const CALLBACK_PATH = "/callback";
 const SCOPES = "org:create_api_key user:profile user:inference";
 
-function parseAuthCode(input: string): { code: string; state?: string } {
-	const trimmed = input.trim();
-	if (!trimmed) return { code: "" };
-
-	try {
-		const url = new URL(trimmed);
-		const code = url.searchParams.get("code") ?? "";
-		const state = url.searchParams.get("state") ?? undefined;
-		if (code) return { code, state };
-	} catch {
-		// Ignore invalid URL parsing and fall back to manual parsing.
-	}
+class AnthropicOAuthFlow extends OAuthCallbackFlow {
+	private verifier: string = "";
+	private challenge: string = "";
 
-	if (trimmed.includes("code=")) {
-		const params = new URLSearchParams(trimmed.replace(/^[?#]/, ""));
-		const code = params.get("code") ?? "";
-		const state = params.get("state") ?? undefined;
-		if (code) return { code, state };
+	constructor(ctrl: OAuthController) {
+		super(ctrl, CALLBACK_PORT, CALLBACK_PATH);
 	}
 
-	const [code, state] = trimmed.split("#");
-	return { code, state };
-}
-
-/**
- * Login with Anthropic OAuth (device code flow)
- *
- * @param onAuthUrl - Callback to handle the authorization URL (e.g., open browser)
- * @param onPromptCode - Callback to prompt user for the authorization code
- */
-export async function loginAnthropic(
-	onAuthUrl: (url: string) => void,
-	onPromptCode: () => Promise<string>,
-): Promise<OAuthCredentials> {
-	const { verifier, challenge } = await generatePKCE();
-
-	// Build authorization URL
-	const authParams = new URLSearchParams({
-		code: "true",
-		client_id: CLIENT_ID,
-		response_type: "code",
-		redirect_uri: REDIRECT_URI,
-		scope: SCOPES,
-		code_challenge: challenge,
-		code_challenge_method: "S256",
-		state: verifier,
-	});
-
-	const authUrl = `${AUTHORIZE_URL}?${authParams.toString()}`;
+	protected async generateAuthUrl(
+		state: string,
+		redirectUri: string,
+	): Promise<{ url: string; instructions?: string }> {
+		const pkce = await generatePKCE();
+		this.verifier = pkce.verifier;
+		this.challenge = pkce.challenge;
 
-	// Notify caller with URL to open
-	onAuthUrl(authUrl);
-
-	// Wait for user to paste authorization code (format: code#state)
-	const authCode = await onPromptCode();
-	const { code, state } = parseAuthCode(authCode);
-
-	// Exchange code for tokens
-	const tokenResponse = await fetch(TOKEN_URL, {
-		method: "POST",
-		headers: {
-			"Content-Type": "application/json",
-		},
-		body: JSON.stringify({
-			grant_type: "authorization_code",
+		const authParams = new URLSearchParams({
+			code: "true",
 			client_id: CLIENT_ID,
-			code,
-			...(state ? { state } : {}),
-			redirect_uri: REDIRECT_URI,
-			code_verifier: verifier,
-		}),
-	});
-
-	if (!tokenResponse.ok) {
-		const error = await tokenResponse.text();
-		throw new Error(`Token exchange failed: ${error}`);
+			response_type: "code",
+			redirect_uri: redirectUri,
+			scope: SCOPES,
+			code_challenge: this.challenge,
+			code_challenge_method: "S256",
+			state,
+		});
+
+		const url = `${AUTHORIZE_URL}?${authParams.toString()}`;
+		return { url };
 	}
 
-	const tokenData = (await tokenResponse.json()) as {
-		access_token: string;
-		refresh_token: string;
-		expires_in: number;
-	};
-
-	// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
-	const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
+	protected async exchangeToken(code: string, state: string, redirectUri: string): Promise<OAuthCredentials> {
+		const tokenResponse = await fetch(TOKEN_URL, {
+			method: "POST",
+			headers: {
+				"Content-Type": "application/json",
+				Accept: "application/json",
+			},
+			body: JSON.stringify({
+				grant_type: "authorization_code",
+				client_id: CLIENT_ID,
+				code,
+				state,
+				redirect_uri: redirectUri,
+				code_verifier: this.verifier,
+			}),
+		});
+
+		if (!tokenResponse.ok) {
+			const error = await tokenResponse.text();
+			throw new Error(`Token exchange failed: ${error}`);
+		}
+
+		const tokenData = (await tokenResponse.json()) as {
+			access_token: string;
+			refresh_token: string;
+			expires_in: number;
+		};
+
+		return {
+			refresh: tokenData.refresh_token,
+			access: tokenData.access_token,
+			expires: Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000,
+		};
+	}
+}
 
-	// Save credentials
-	return {
-		refresh: tokenData.refresh_token,
-		access: tokenData.access_token,
-		expires: expiresAt,
-	};
+/**
+ * Login with Anthropic OAuth
+ */
+export async function loginAnthropic(ctrl: OAuthController): Promise<OAuthCredentials> {
+	const flow = new AnthropicOAuthFlow(ctrl);
+	return flow.login();
 }
 
 /**
@@ -113,7 +95,7 @@ export async function loginAnthropic(
 export async function refreshAnthropicToken(refreshToken: string): Promise<OAuthCredentials> {
 	const response = await fetch(TOKEN_URL, {
 		method: "POST",
-		headers: { "Content-Type": "application/json" },
+		headers: { "Content-Type": "application/json", Accept: "application/json" },
 		body: JSON.stringify({
 			grant_type: "refresh_token",
 			client_id: CLIENT_ID,

package/src/utils/oauth/callback-server.ts

@@ -0,0 +1,245 @@
+/**
+ * Abstract base class for OAuth flows with local callback servers.
+ *
+ * Handles:
+ * - Port allocation (tries expected port, falls back to random)
+ * - Callback server setup and request handling
+ * - Common OAuth flow logic
+ *
+ * Providers extend this and implement:
+ * - generateAuthUrl(): Build provider-specific authorization URL
+ * - exchangeToken(): Exchange authorization code for tokens
+ */
+
+import templateHtml from "./oauth.html" with { type: "text" };
+import type { OAuthController, OAuthCredentials } from "./types";
+
+const DEFAULT_TIMEOUT = 120;
+const DEFAULT_HOSTNAME = "localhost";
+const CALLBACK_PATH = "/callback";
+
+export type CallbackResult = { code: string; state: string };
+
+/**
+ * Abstract base class for OAuth flows with local callback servers.
+ */
+export abstract class OAuthCallbackFlow {
+	protected ctrl: OAuthController;
+	protected preferredPort: number;
+	protected callbackPath: string;
+	private callbackResolve?: (result: CallbackResult) => void;
+	private callbackReject?: (error: string) => void;
+
+	constructor(ctrl: OAuthController, preferredPort: number, callbackPath: string = CALLBACK_PATH) {
+		this.ctrl = ctrl;
+		this.preferredPort = preferredPort;
+		this.callbackPath = callbackPath;
+	}
+
+	/**
+	 * Generate provider-specific authorization URL.
+	 * @param state - CSRF state token
+	 * @param redirectUri - The actual redirect URI to use (may differ from expected if port fallback occurred)
+	 * @returns Authorization URL and optional instructions
+	 */
+	protected abstract generateAuthUrl(
+		state: string,
+		redirectUri: string,
+	): Promise<{ url: string; instructions?: string }>;
+
+	/**
+	 * Exchange authorization code for OAuth tokens.
+	 * @param code - Authorization code from callback
+	 * @param state - CSRF state token
+	 * @param redirectUri - The actual redirect URI used (must match authorization request)
+	 * @returns OAuth credentials
+	 */
+	protected abstract exchangeToken(code: string, state: string, redirectUri: string): Promise<OAuthCredentials>;
+
+	/**
+	 * Generate CSRF state token. Override if provider needs custom state generation.
+	 */
+	protected generateState(): string {
+		const bytes = new Uint8Array(16);
+		crypto.getRandomValues(bytes);
+		return Array.from(bytes)
+			.map((value) => value.toString(16).padStart(2, "0"))
+			.join("");
+	}
+
+	/**
+	 * Execute the OAuth login flow.
+	 */
+	async login(): Promise<OAuthCredentials> {
+		const state = this.generateState();
+
+		// Start callback server first to get actual redirect URI
+		const { server, redirectUri } = await this.startCallbackServer(state);
+
+		try {
+			// Generate auth URL with the ACTUAL redirect URI (may differ from expected if port was busy)
+			const { url: authUrl, instructions } = await this.generateAuthUrl(state, redirectUri);
+
+			// Notify controller that auth is ready
+			this.ctrl.onAuth?.({ url: authUrl, instructions });
+			this.ctrl.onProgress?.("Waiting for browser authentication...");
+
+			// Wait for callback or manual input
+			const { code } = await this.waitForCallback(state);
+
+			this.ctrl.onProgress?.("Exchanging authorization code for tokens...");
+
+			return await this.exchangeToken(code, state, redirectUri);
+		} finally {
+			server.stop();
+		}
+	}
+
+	/**
+	 * Start callback server, trying preferred port first, falling back to random.
+	 */
+	private async startCallbackServer(
+		expectedState: string,
+	): Promise<{ server: Bun.Server<unknown>; redirectUri: string }> {
+		// Try preferred port first
+		try {
+			const redirectUri = `http://${DEFAULT_HOSTNAME}:${this.preferredPort}${this.callbackPath}`;
+			const server = this.createServer(this.preferredPort, expectedState);
+			return { server, redirectUri };
+		} catch {
+			// Port busy or unavailable, try random port
+			const randomPort = 0; // Let OS assign
+			const server = this.createServer(randomPort, expectedState);
+			const actualPort = server.port;
+			const redirectUri = `http://${DEFAULT_HOSTNAME}:${actualPort}${this.callbackPath}`;
+			this.ctrl.onProgress?.(`Preferred port ${this.preferredPort} unavailable, using port ${actualPort}`);
+			return { server, redirectUri };
+		}
+	}
+
+	/**
+	 * Create HTTP server for OAuth callback.
+	 */
+	private createServer(port: number, expectedState: string): Bun.Server<unknown> {
+		return Bun.serve({
+			hostname: DEFAULT_HOSTNAME,
+			port,
+			reusePort: false,
+			fetch: (req) => this.handleCallback(req, expectedState),
+		});
+	}
+
+	/**
+	 * Handle OAuth callback HTTP request.
+	 */
+	private handleCallback(req: Request, expectedState: string): Response {
+		const url = new URL(req.url);
+
+		if (url.pathname !== this.callbackPath) {
+			return new Response("Not Found", { status: 404 });
+		}
+
+		const code = url.searchParams.get("code");
+		const state = url.searchParams.get("state") || "";
+		const error = url.searchParams.get("error") || "";
+		const errorDescription = url.searchParams.get("error_description") || error;
+
+		type OkState = { ok: true; code: string; state: string };
+		type ErrorState = { ok?: false; error?: string };
+		let resultState: OkState | ErrorState;
+
+		if (error) {
+			resultState = { ok: false, error: `Authorization failed: ${errorDescription}` };
+		} else if (!code) {
+			resultState = { ok: false, error: "Missing authorization code" };
+		} else if (expectedState && state !== expectedState) {
+			resultState = { ok: false, error: "State mismatch - possible CSRF attack" };
+		} else {
+			resultState = { ok: true, code, state };
+		}
+
+		// Signal to waitForCallback
+		queueMicrotask(() => {
+			if (resultState.ok) {
+				this.callbackResolve?.({ code: resultState.code, state: resultState.state });
+			} else {
+				this.callbackReject?.(resultState.error ?? "Unknown error");
+			}
+		});
+
+		return new Response(
+			(templateHtml as unknown as string).replaceAll("__OAUTH_STATE__", JSON.stringify(resultState)),
+			{
+				status: resultState.ok ? 200 : 500,
+				headers: { "Content-Type": "text/html" },
+			},
+		);
+	}
+
+	/**
+	 * Wait for OAuth callback or manual input (whichever comes first).
+	 */
+	private waitForCallback(expectedState: string): Promise<CallbackResult> {
+		const timeoutSignal = AbortSignal.timeout(DEFAULT_TIMEOUT * 1000);
+		const signal = this.ctrl.signal ? AbortSignal.any([this.ctrl.signal, timeoutSignal]) : timeoutSignal;
+
+		const callbackPromise = new Promise<CallbackResult>((resolve, reject) => {
+			this.callbackResolve = resolve;
+			this.callbackReject = reject;
+
+			signal.addEventListener("abort", () => {
+				this.callbackResolve = undefined;
+				this.callbackReject = undefined;
+				reject(new Error(`OAuth callback cancelled: ${signal.reason}`));
+			});
+		});
+
+		// Manual input race (if supported)
+		if (this.ctrl.onManualCodeInput) {
+			const manualPromise = this.ctrl.onManualCodeInput().then((input): CallbackResult => {
+				const parsed = parseCallbackInput(input);
+				if (!parsed.code) {
+					throw new Error("No authorization code found in input");
+				}
+				if (expectedState && parsed.state && parsed.state !== expectedState) {
+					throw new Error("State mismatch - possible CSRF attack");
+				}
+				return { code: parsed.code, state: parsed.state ?? "" };
+			});
+
+			return Promise.race([callbackPromise, manualPromise]);
+		}
+
+		return callbackPromise;
+	}
+}
+
+/**
+ * Parse a redirect URL or code string to extract code and state.
+ */
+export function parseCallbackInput(input: string): { code?: string; state?: string } {
+	const value = input.trim();
+	if (!value) return {};
+
+	try {
+		const url = new URL(value);
+		return {
+			code: url.searchParams.get("code") ?? undefined,
+			state: url.searchParams.get("state") ?? undefined,
+		};
+	} catch {
+		// Not a URL - check for query string format
+	}
+
+	if (value.includes("code=")) {
+		const params = new URLSearchParams(value.replace(/^[?#]/, ""));
+		return {
+			code: params.get("code") ?? undefined,
+			state: params.get("state") ?? undefined,
+		};
+	}
+
+	// Assume raw code, possibly with state after #
+	const [code, state] = value.split("#", 2);
+	return { code, state };
+}

package/src/utils/oauth/cursor.ts

@@ -10,10 +10,6 @@ const POLL_BASE_DELAY = 1000;
 const POLL_MAX_DELAY = 10000;
 const POLL_BACKOFF_MULTIPLIER = 1.2;
 
-function sleep(ms: number): Promise<void> {
-	return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 export interface CursorAuthParams {
 	verifier: string;
 	challenge: string;
@@ -45,7 +41,7 @@ export async function pollCursorAuth(
 	let consecutiveErrors = 0;
 
 	for (let attempt = 0; attempt < POLL_MAX_ATTEMPTS; attempt++) {
-		await sleep(delay);
+		await Bun.sleep(delay);
 
 		try {
 			const response = await fetch(`${CURSOR_POLL_URL}?uuid=${uuid}&verifier=${verifier}`);

package/src/utils/oauth/github-copilot.ts

@@ -2,6 +2,7 @@
  * GitHub Copilot OAuth flow
  */
 
+import { abortableSleep } from "@oh-my-pi/pi-utils";
 import { getModels } from "../../models";
 import type { OAuthCredentials } from "./types";
 
@@ -136,29 +137,6 @@ async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
 	};
 }
 
-/**
- * Sleep that can be interrupted by an AbortSignal
- */
-function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
-	return new Promise((resolve, reject) => {
-		if (signal?.aborted) {
-			reject(new Error("Login cancelled"));
-			return;
-		}
-
-		const timeout = setTimeout(resolve, ms);
-
-		signal?.addEventListener(
-			"abort",
-			() => {
-				clearTimeout(timeout);
-				reject(new Error("Login cancelled"));
-			},
-			{ once: true },
-		);
-	});
-}
-
 async function pollForGitHubAccessToken(
 	domain: string,
 	deviceCode: string,