@lenylvt/pi-coding-agent 0.64.0

package/src/core/auth-storage.ts

@@ -0,0 +1,488 @@
+/**
+ * Credential storage for API keys and OAuth tokens.
+ * Handles loading, saving, and refreshing credentials from auth.json.
+ *
+ * Uses file locking to prevent race conditions when multiple pi instances
+ * try to refresh tokens simultaneously.
+ */
+
+import { getEnvApiKey, type OAuthCredentials, type OAuthLoginCallbacks, type OAuthProviderId } from "@lenylvt/pi-ai";
+import { getOAuthApiKey, getOAuthProvider, getOAuthProviders } from "@lenylvt/pi-ai/oauth";
+import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
+import { dirname, join } from "path";
+import lockfile from "proper-lockfile";
+import { getAgentDir } from "../config.js";
+import { resolveConfigValue } from "./resolve-config-value.js";
+
+export type ApiKeyCredential = {
+	type: "api_key";
+	key: string;
+};
+
+export type OAuthCredential = {
+	type: "oauth";
+} & OAuthCredentials;
+
+export type AuthCredential = ApiKeyCredential | OAuthCredential;
+
+export type AuthStorageData = Record<string, AuthCredential>;
+
+type LockResult<T> = {
+	result: T;
+	next?: string;
+};
+
+export interface AuthStorageBackend {
+	withLock<T>(fn: (current: string | undefined) => LockResult<T>): T;
+	withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T>;
+}
+
+export class FileAuthStorageBackend implements AuthStorageBackend {
+	constructor(private authPath: string = join(getAgentDir(), "auth.json")) {}
+
+	private ensureParentDir(): void {
+		const dir = dirname(this.authPath);
+		if (!existsSync(dir)) {
+			mkdirSync(dir, { recursive: true, mode: 0o700 });
+		}
+	}
+
+	private ensureFileExists(): void {
+		if (!existsSync(this.authPath)) {
+			writeFileSync(this.authPath, "{}", "utf-8");
+			chmodSync(this.authPath, 0o600);
+		}
+	}
+
+	private acquireLockSyncWithRetry(path: string): () => void {
+		const maxAttempts = 10;
+		const delayMs = 20;
+		let lastError: unknown;
+
+		for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+			try {
+				return lockfile.lockSync(path, { realpath: false });
+			} catch (error) {
+				const code =
+					typeof error === "object" && error !== null && "code" in error
+						? String((error as { code?: unknown }).code)
+						: undefined;
+				if (code !== "ELOCKED" || attempt === maxAttempts) {
+					throw error;
+				}
+				lastError = error;
+				const start = Date.now();
+				while (Date.now() - start < delayMs) {
+					// Sleep synchronously to avoid changing callers to async.
+				}
+			}
+		}
+
+		throw (lastError as Error) ?? new Error("Failed to acquire auth storage lock");
+	}
+
+	withLock<T>(fn: (current: string | undefined) => LockResult<T>): T {
+		this.ensureParentDir();
+		this.ensureFileExists();
+
+		let release: (() => void) | undefined;
+		try {
+			release = this.acquireLockSyncWithRetry(this.authPath);
+			const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined;
+			const { result, next } = fn(current);
+			if (next !== undefined) {
+				writeFileSync(this.authPath, next, "utf-8");
+				chmodSync(this.authPath, 0o600);
+			}
+			return result;
+		} finally {
+			if (release) {
+				release();
+			}
+		}
+	}
+
+	async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> {
+		this.ensureParentDir();
+		this.ensureFileExists();
+
+		let release: (() => Promise<void>) | undefined;
+		let lockCompromised = false;
+		let lockCompromisedError: Error | undefined;
+		const throwIfCompromised = () => {
+			if (lockCompromised) {
+				throw lockCompromisedError ?? new Error("Auth storage lock was compromised");
+			}
+		};
+
+		try {
+			release = await lockfile.lock(this.authPath, {
+				retries: {
+					retries: 10,
+					factor: 2,
+					minTimeout: 100,
+					maxTimeout: 10000,
+					randomize: true,
+				},
+				stale: 30000,
+				onCompromised: (err) => {
+					lockCompromised = true;
+					lockCompromisedError = err;
+				},
+			});
+
+			throwIfCompromised();
+			const current = existsSync(this.authPath) ? readFileSync(this.authPath, "utf-8") : undefined;
+			const { result, next } = await fn(current);
+			throwIfCompromised();
+			if (next !== undefined) {
+				writeFileSync(this.authPath, next, "utf-8");
+				chmodSync(this.authPath, 0o600);
+			}
+			throwIfCompromised();
+			return result;
+		} finally {
+			if (release) {
+				try {
+					await release();
+				} catch {
+					// Ignore unlock errors when lock is compromised.
+				}
+			}
+		}
+	}
+}
+
+export class InMemoryAuthStorageBackend implements AuthStorageBackend {
+	private value: string | undefined;
+
+	withLock<T>(fn: (current: string | undefined) => LockResult<T>): T {
+		const { result, next } = fn(this.value);
+		if (next !== undefined) {
+			this.value = next;
+		}
+		return result;
+	}
+
+	async withLockAsync<T>(fn: (current: string | undefined) => Promise<LockResult<T>>): Promise<T> {
+		const { result, next } = await fn(this.value);
+		if (next !== undefined) {
+			this.value = next;
+		}
+		return result;
+	}
+}
+
+/**
+ * Credential storage backed by a JSON file.
+ */
+export class AuthStorage {
+	private data: AuthStorageData = {};
+	private runtimeOverrides: Map<string, string> = new Map();
+	private fallbackResolver?: (provider: string) => string | undefined;
+	private loadError: Error | null = null;
+	private errors: Error[] = [];
+
+	private constructor(private storage: AuthStorageBackend) {
+		this.reload();
+	}
+
+	static create(authPath?: string): AuthStorage {
+		return new AuthStorage(new FileAuthStorageBackend(authPath ?? join(getAgentDir(), "auth.json")));
+	}
+
+	static fromStorage(storage: AuthStorageBackend): AuthStorage {
+		return new AuthStorage(storage);
+	}
+
+	static inMemory(data: AuthStorageData = {}): AuthStorage {
+		const storage = new InMemoryAuthStorageBackend();
+		storage.withLock(() => ({ result: undefined, next: JSON.stringify(data, null, 2) }));
+		return AuthStorage.fromStorage(storage);
+	}
+
+	/**
+	 * Set a runtime API key override (not persisted to disk).
+	 * Used for CLI --api-key flag.
+	 */
+	setRuntimeApiKey(provider: string, apiKey: string): void {
+		this.runtimeOverrides.set(provider, apiKey);
+	}
+
+	/**
+	 * Remove a runtime API key override.
+	 */
+	removeRuntimeApiKey(provider: string): void {
+		this.runtimeOverrides.delete(provider);
+	}
+
+	/**
+	 * Set a fallback resolver for API keys not found in auth.json or env vars.
+	 * Used for custom provider keys from models.json.
+	 */
+	setFallbackResolver(resolver: (provider: string) => string | undefined): void {
+		this.fallbackResolver = resolver;
+	}
+
+	private recordError(error: unknown): void {
+		const normalizedError = error instanceof Error ? error : new Error(String(error));
+		this.errors.push(normalizedError);
+	}
+
+	private parseStorageData(content: string | undefined): AuthStorageData {
+		if (!content) {
+			return {};
+		}
+		return JSON.parse(content) as AuthStorageData;
+	}
+
+	/**
+	 * Reload credentials from storage.
+	 */
+	reload(): void {
+		let content: string | undefined;
+		try {
+			this.storage.withLock((current) => {
+				content = current;
+				return { result: undefined };
+			});
+			this.data = this.parseStorageData(content);
+			this.loadError = null;
+		} catch (error) {
+			this.loadError = error as Error;
+			this.recordError(error);
+		}
+	}
+
+	private persistProviderChange(provider: string, credential: AuthCredential | undefined): void {
+		if (this.loadError) {
+			return;
+		}
+
+		try {
+			this.storage.withLock((current) => {
+				const currentData = this.parseStorageData(current);
+				const merged: AuthStorageData = { ...currentData };
+				if (credential) {
+					merged[provider] = credential;
+				} else {
+					delete merged[provider];
+				}
+				return { result: undefined, next: JSON.stringify(merged, null, 2) };
+			});
+		} catch (error) {
+			this.recordError(error);
+		}
+	}
+
+	/**
+	 * Get credential for a provider.
+	 */
+	get(provider: string): AuthCredential | undefined {
+		return this.data[provider] ?? undefined;
+	}
+
+	/**
+	 * Set credential for a provider.
+	 */
+	set(provider: string, credential: AuthCredential): void {
+		this.data[provider] = credential;
+		this.persistProviderChange(provider, credential);
+	}
+
+	/**
+	 * Remove credential for a provider.
+	 */
+	remove(provider: string): void {
+		delete this.data[provider];
+		this.persistProviderChange(provider, undefined);
+	}
+
+	/**
+	 * List all providers with credentials.
+	 */
+	list(): string[] {
+		return Object.keys(this.data);
+	}
+
+	/**
+	 * Check if credentials exist for a provider in auth.json.
+	 */
+	has(provider: string): boolean {
+		return provider in this.data;
+	}
+
+	/**
+	 * Check if any form of auth is configured for a provider.
+	 * Unlike getApiKey(), this doesn't refresh OAuth tokens.
+	 */
+	hasAuth(provider: string): boolean {
+		if (this.runtimeOverrides.has(provider)) return true;
+		if (this.data[provider]) return true;
+		if (getEnvApiKey(provider)) return true;
+		if (this.fallbackResolver?.(provider)) return true;
+		return false;
+	}
+
+	/**
+	 * Get all credentials (for passing to getOAuthApiKey).
+	 */
+	getAll(): AuthStorageData {
+		return { ...this.data };
+	}
+
+	drainErrors(): Error[] {
+		const drained = [...this.errors];
+		this.errors = [];
+		return drained;
+	}
+
+	/**
+	 * Login to an OAuth provider.
+	 */
+	async login(providerId: OAuthProviderId, callbacks: OAuthLoginCallbacks): Promise<void> {
+		const provider = getOAuthProvider(providerId);
+		if (!provider) {
+			throw new Error(`Unknown OAuth provider: ${providerId}`);
+		}
+
+		const credentials = await provider.login(callbacks);
+		this.set(providerId, { type: "oauth", ...credentials });
+	}
+
+	/**
+	 * Logout from a provider.
+	 */
+	logout(provider: string): void {
+		this.remove(provider);
+	}
+
+	/**
+	 * Refresh OAuth token with backend locking to prevent race conditions.
+	 * Multiple pi instances may try to refresh simultaneously when tokens expire.
+	 */
+	private async refreshOAuthTokenWithLock(
+		providerId: OAuthProviderId,
+	): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> {
+		const provider = getOAuthProvider(providerId);
+		if (!provider) {
+			return null;
+		}
+
+		const result = await this.storage.withLockAsync(async (current) => {
+			const currentData = this.parseStorageData(current);
+			this.data = currentData;
+			this.loadError = null;
+
+			const cred = currentData[providerId];
+			if (cred?.type !== "oauth") {
+				return { result: null };
+			}
+
+			if (Date.now() < cred.expires) {
+				return { result: { apiKey: provider.getApiKey(cred), newCredentials: cred } };
+			}
+
+			const oauthCreds: Record<string, OAuthCredentials> = {};
+			for (const [key, value] of Object.entries(currentData)) {
+				if (value.type === "oauth") {
+					oauthCreds[key] = value;
+				}
+			}
+
+			const refreshed = await getOAuthApiKey(providerId, oauthCreds);
+			if (!refreshed) {
+				return { result: null };
+			}
+
+			const merged: AuthStorageData = {
+				...currentData,
+				[providerId]: { type: "oauth", ...refreshed.newCredentials },
+			};
+			this.data = merged;
+			this.loadError = null;
+			return { result: refreshed, next: JSON.stringify(merged, null, 2) };
+		});
+
+		return result;
+	}
+
+	/**
+	 * Get API key for a provider.
+	 * Priority:
+	 * 1. Runtime override (CLI --api-key)
+	 * 2. API key from auth.json
+	 * 3. OAuth token from auth.json (auto-refreshed with locking)
+	 * 4. Environment variable
+	 * 5. Fallback resolver (models.json custom providers)
+	 */
+	async getApiKey(providerId: string, options?: { includeFallback?: boolean }): Promise<string | undefined> {
+		// Runtime override takes highest priority
+		const runtimeKey = this.runtimeOverrides.get(providerId);
+		if (runtimeKey) {
+			return runtimeKey;
+		}
+
+		const cred = this.data[providerId];
+
+		if (cred?.type === "api_key") {
+			return resolveConfigValue(cred.key);
+		}
+
+		if (cred?.type === "oauth") {
+			const provider = getOAuthProvider(providerId);
+			if (!provider) {
+				// Unknown OAuth provider, can't get API key
+				return undefined;
+			}
+
+			// Check if token needs refresh
+			const needsRefresh = Date.now() >= cred.expires;
+
+			if (needsRefresh) {
+				// Use locked refresh to prevent race conditions
+				try {
+					const result = await this.refreshOAuthTokenWithLock(providerId);
+					if (result) {
+						return result.apiKey;
+					}
+				} catch (error) {
+					this.recordError(error);
+					// Refresh failed - re-read file to check if another instance succeeded
+					this.reload();
+					const updatedCred = this.data[providerId];
+
+					if (updatedCred?.type === "oauth" && Date.now() < updatedCred.expires) {
+						// Another instance refreshed successfully, use those credentials
+						return provider.getApiKey(updatedCred);
+					}
+
+					// Refresh truly failed - return undefined so model discovery skips this provider
+					// User can /login to re-authenticate (credentials preserved for retry)
+					return undefined;
+				}
+			} else {
+				// Token not expired, use current access token
+				return provider.getApiKey(cred);
+			}
+		}
+
+		// Fall back to environment variable
+		const envKey = getEnvApiKey(providerId);
+		if (envKey) return envKey;
+
+		// Fall back to custom resolver (e.g., models.json custom providers)
+		if (options?.includeFallback !== false) {
+			return this.fallbackResolver?.(providerId) ?? undefined;
+		}
+
+		return undefined;
+	}
+
+	/**
+	 * Get all registered OAuth providers
+	 */
+	getOAuthProviders() {
+		return getOAuthProviders();
+	}
+}

package/src/core/bash-executor.ts

@@ -0,0 +1,158 @@
+/**
+ * Bash command execution with streaming support and cancellation.
+ *
+ * This module provides a unified bash execution implementation used by:
+ * - AgentSession.executeBash() for interactive and RPC modes
+ * - Direct calls from modes that need bash execution
+ */
+
+import { randomBytes } from "node:crypto";
+import { createWriteStream, type WriteStream } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import stripAnsi from "strip-ansi";
+import { sanitizeBinaryOutput } from "../utils/shell.js";
+import { type BashOperations, createLocalBashOperations } from "./tools/bash.js";
+import { DEFAULT_MAX_BYTES, truncateTail } from "./tools/truncate.js";
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface BashExecutorOptions {
+	/** Callback for streaming output chunks (already sanitized) */
+	onChunk?: (chunk: string) => void;
+	/** AbortSignal for cancellation */
+	signal?: AbortSignal;
+}
+
+export interface BashResult {
+	/** Combined stdout + stderr output (sanitized, possibly truncated) */
+	output: string;
+	/** Process exit code (undefined if killed/cancelled) */
+	exitCode: number | undefined;
+	/** Whether the command was cancelled via signal */
+	cancelled: boolean;
+	/** Whether the output was truncated */
+	truncated: boolean;
+	/** Path to temp file containing full output (if output exceeded truncation threshold) */
+	fullOutputPath?: string;
+}
+
+// ============================================================================
+// Implementation
+// ============================================================================
+
+/**
+ * Execute a bash command with optional streaming and cancellation support.
+ *
+ * Uses the same local BashOperations backend as createBashTool() so interactive
+ * user bash and tool-invoked bash share the same process spawning behavior.
+ * Sanitization, newline normalization, temp-file capture, and truncation still
+ * happen in executeBashWithOperations(), so reusing the local backend does not
+ * change output processing behavior.
+ *
+ * @param command - The bash command to execute
+ * @param options - Optional streaming callback and abort signal
+ * @returns Promise resolving to execution result
+ */
+export function executeBash(command: string, options?: BashExecutorOptions): Promise<BashResult> {
+	return executeBashWithOperations(command, process.cwd(), createLocalBashOperations(), options);
+}
+
+/**
+ * Execute a bash command using custom BashOperations.
+ * Used for remote execution (SSH, containers, etc.).
+ */
+export async function executeBashWithOperations(
+	command: string,
+	cwd: string,
+	operations: BashOperations,
+	options?: BashExecutorOptions,
+): Promise<BashResult> {
+	const outputChunks: string[] = [];
+	let outputBytes = 0;
+	const maxOutputBytes = DEFAULT_MAX_BYTES * 2;
+
+	let tempFilePath: string | undefined;
+	let tempFileStream: WriteStream | undefined;
+	let totalBytes = 0;
+
+	const decoder = new TextDecoder();
+
+	const onData = (data: Buffer) => {
+		totalBytes += data.length;
+
+		// Sanitize: strip ANSI, replace binary garbage, normalize newlines
+		const text = sanitizeBinaryOutput(stripAnsi(decoder.decode(data, { stream: true }))).replace(/\r/g, "");
+
+		// Start writing to temp file if exceeds threshold
+		if (totalBytes > DEFAULT_MAX_BYTES && !tempFilePath) {
+			const id = randomBytes(8).toString("hex");
+			tempFilePath = join(tmpdir(), `pi-bash-${id}.log`);
+			tempFileStream = createWriteStream(tempFilePath);
+			for (const chunk of outputChunks) {
+				tempFileStream.write(chunk);
+			}
+		}
+
+		if (tempFileStream) {
+			tempFileStream.write(text);
+		}
+
+		// Keep rolling buffer
+		outputChunks.push(text);
+		outputBytes += text.length;
+		while (outputBytes > maxOutputBytes && outputChunks.length > 1) {
+			const removed = outputChunks.shift()!;
+			outputBytes -= removed.length;
+		}
+
+		// Stream to callback
+		if (options?.onChunk) {
+			options.onChunk(text);
+		}
+	};
+
+	try {
+		const result = await operations.exec(command, cwd, {
+			onData,
+			signal: options?.signal,
+		});
+
+		if (tempFileStream) {
+			tempFileStream.end();
+		}
+
+		const fullOutput = outputChunks.join("");
+		const truncationResult = truncateTail(fullOutput);
+		const cancelled = options?.signal?.aborted ?? false;
+
+		return {
+			output: truncationResult.truncated ? truncationResult.content : fullOutput,
+			exitCode: cancelled ? undefined : (result.exitCode ?? undefined),
+			cancelled,
+			truncated: truncationResult.truncated,
+			fullOutputPath: tempFilePath,
+		};
+	} catch (err) {
+		if (tempFileStream) {
+			tempFileStream.end();
+		}
+
+		// Check if it was an abort
+		if (options?.signal?.aborted) {
+			const fullOutput = outputChunks.join("");
+			const truncationResult = truncateTail(fullOutput);
+			return {
+				output: truncationResult.truncated ? truncationResult.content : fullOutput,
+				exitCode: undefined,
+				cancelled: true,
+				truncated: truncationResult.truncated,
+				fullOutputPath: tempFilePath,
+			};
+		}
+
+		throw err;
+	}
+}