@oh-my-pi/pi-coding-agent 0.1.0

package/src/core/auth-storage.ts

@@ -0,0 +1,258 @@
+/**
+ * Credential storage for API keys and OAuth tokens.
+ * Handles loading, saving, and refreshing credentials from auth.json.
+ */
+
+import { chmodSync, existsSync, mkdirSync, readFileSync } from "node:fs";
+import { dirname } from "node:path";
+import {
+	getEnvApiKey,
+	getOAuthApiKey,
+	loginAnthropic,
+	loginAntigravity,
+	loginGeminiCli,
+	loginGitHubCopilot,
+	type OAuthCredentials,
+	type OAuthProvider,
+} from "@oh-my-pi/pi-ai";
+import { logger } from "./logger";
+
+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>;
+
+/**
+ * Credential storage backed by a JSON file.
+ * Reads from multiple fallback paths, writes to primary path.
+ */
+export class AuthStorage {
+	private data: AuthStorageData = {};
+	private runtimeOverrides: Map<string, string> = new Map();
+	private fallbackResolver?: (provider: string) => string | undefined;
+
+	/**
+	 * @param authPath - Primary path for reading/writing auth.json
+	 * @param fallbackPaths - Additional paths to check when reading (legacy support)
+	 */
+	constructor(
+		private authPath: string,
+		private fallbackPaths: string[] = [],
+	) {
+		this.reload();
+	}
+
+	/**
+	 * 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;
+	}
+
+	/**
+	 * Reload credentials from disk.
+	 * Checks primary path first, then fallback paths.
+	 */
+	reload(): void {
+		const pathsToCheck = [this.authPath, ...this.fallbackPaths];
+
+		logger.debug("AuthStorage.reload checking paths", { paths: pathsToCheck });
+
+		for (const authPath of pathsToCheck) {
+			const exists = existsSync(authPath);
+			logger.debug("AuthStorage.reload path check", { path: authPath, exists });
+
+			if (exists) {
+				try {
+					this.data = JSON.parse(readFileSync(authPath, "utf-8"));
+					logger.debug("AuthStorage.reload loaded", { path: authPath, providers: Object.keys(this.data) });
+					return;
+				} catch (e) {
+					logger.error("AuthStorage failed to parse auth file", { path: authPath, error: String(e) });
+					// Continue to next path on parse error
+				}
+			}
+		}
+
+		logger.warn("AuthStorage no auth file found", { checkedPaths: pathsToCheck });
+		this.data = {};
+	}
+
+	/**
+	 * Save credentials to disk.
+	 */
+	private save(): void {
+		const dir = dirname(this.authPath);
+		if (!existsSync(dir)) {
+			mkdirSync(dir, { recursive: true, mode: 0o700 });
+		}
+		Bun.write(this.authPath, JSON.stringify(this.data, null, 2));
+		chmodSync(this.authPath, 0o600);
+	}
+
+	/**
+	 * 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.save();
+	}
+
+	/**
+	 * Remove credential for a provider.
+	 */
+	remove(provider: string): void {
+		delete this.data[provider];
+		this.save();
+	}
+
+	/**
+	 * List all providers with credentials.
+	 */
+	list(): string[] {
+		return Object.keys(this.data);
+	}
+
+	/**
+	 * Check if credentials exist for a provider.
+	 */
+	has(provider: string): boolean {
+		return provider in this.data;
+	}
+
+	/**
+	 * Get all credentials (for passing to getOAuthApiKey).
+	 */
+	getAll(): AuthStorageData {
+		return { ...this.data };
+	}
+
+	/**
+	 * Login to an OAuth provider.
+	 */
+	async login(
+		provider: OAuthProvider,
+		callbacks: {
+			onAuth: (info: { url: string; instructions?: string }) => void;
+			onPrompt: (prompt: { message: string; placeholder?: string }) => Promise<string>;
+			onProgress?: (message: string) => void;
+		},
+	): Promise<void> {
+		let credentials: OAuthCredentials;
+
+		switch (provider) {
+			case "anthropic":
+				credentials = await loginAnthropic(
+					(url) => callbacks.onAuth({ url }),
+					() => callbacks.onPrompt({ message: "Paste the authorization code:" }),
+				);
+				break;
+			case "github-copilot":
+				credentials = await loginGitHubCopilot({
+					onAuth: (url, instructions) => callbacks.onAuth({ url, instructions }),
+					onPrompt: callbacks.onPrompt,
+					onProgress: callbacks.onProgress,
+				});
+				break;
+			case "google-gemini-cli":
+				credentials = await loginGeminiCli(callbacks.onAuth, callbacks.onProgress);
+				break;
+			case "google-antigravity":
+				credentials = await loginAntigravity(callbacks.onAuth, callbacks.onProgress);
+				break;
+			default:
+				throw new Error(`Unknown OAuth provider: ${provider}`);
+		}
+
+		this.set(provider, { type: "oauth", ...credentials });
+	}
+
+	/**
+	 * Logout from a provider.
+	 */
+	logout(provider: string): void {
+		this.remove(provider);
+	}
+
+	/**
+	 * 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)
+	 * 4. Environment variable
+	 * 5. Fallback resolver (models.json custom providers)
+	 */
+	async getApiKey(provider: string): Promise<string | undefined> {
+		// Runtime override takes highest priority
+		const runtimeKey = this.runtimeOverrides.get(provider);
+		if (runtimeKey) {
+			return runtimeKey;
+		}
+
+		const cred = this.data[provider];
+
+		if (cred?.type === "api_key") {
+			return cred.key;
+		}
+
+		if (cred?.type === "oauth") {
+			// Filter to only oauth credentials for getOAuthApiKey
+			const oauthCreds: Record<string, OAuthCredentials> = {};
+			for (const [key, value] of Object.entries(this.data)) {
+				if (value.type === "oauth") {
+					oauthCreds[key] = value;
+				}
+			}
+
+			try {
+				const result = await getOAuthApiKey(provider as OAuthProvider, oauthCreds);
+				if (result) {
+					this.data[provider] = { type: "oauth", ...result.newCredentials };
+					this.save();
+					return result.apiKey;
+				}
+			} catch {
+				this.remove(provider);
+			}
+		}
+
+		// Fall back to environment variable
+		const envKey = getEnvApiKey(provider);
+		if (envKey) return envKey;
+
+		// Fall back to custom resolver (e.g., models.json custom providers)
+		return this.fallbackResolver?.(provider) ?? undefined;
+	}
+}

package/src/core/bash-executor.ts

@@ -0,0 +1,197 @@
+/**
+ * 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 { createWriteStream, type WriteStream } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { Subprocess } from "bun";
+import stripAnsi from "strip-ansi";
+import { getShellConfig, killProcessTree, sanitizeBinaryOutput } from "../utils/shell";
+import { getOrCreateSnapshot, getSnapshotSourceCommand } from "../utils/shell-snapshot";
+import { DEFAULT_MAX_BYTES, truncateTail } from "./tools/truncate";
+import { ScopeSignal } from "./utils";
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface BashExecutorOptions {
+	/** Working directory for command execution */
+	cwd?: string;
+	/** Timeout in milliseconds */
+	timeout?: number;
+	/** 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
+// ============================================================================
+
+function createSanitizer(): TransformStream<Uint8Array, string> {
+	const decoder = new TextDecoder();
+	return new TransformStream({
+		transform(chunk, controller) {
+			const text = sanitizeBinaryOutput(stripAnsi(decoder.decode(chunk, { stream: true }))).replace(/\r/g, "");
+			controller.enqueue(text);
+		},
+	});
+}
+
+function createOutputSink(
+	spillThreshold: number,
+	maxBuffer: number,
+	onChunk?: (text: string) => void,
+): WritableStream<string> & {
+	dump: (annotation?: string) => { output: string; truncated: boolean; fullOutputPath?: string };
+} {
+	const chunks: string[] = [];
+	let chunkBytes = 0;
+	let totalBytes = 0;
+	let fullOutputPath: string | undefined;
+	let fullOutputStream: WriteStream | undefined;
+
+	const sink = new WritableStream<string>({
+		write(text) {
+			totalBytes += text.length;
+
+			// Spill to temp file if needed
+			if (totalBytes > spillThreshold && !fullOutputPath) {
+				fullOutputPath = join(tmpdir(), `omp-${crypto.randomUUID()}.buffer`);
+				const ts = createWriteStream(fullOutputPath);
+				chunks.forEach((c) => {
+					ts.write(c);
+				});
+				fullOutputStream = ts;
+			}
+			fullOutputStream?.write(text);
+
+			// Rolling buffer
+			chunks.push(text);
+			chunkBytes += text.length;
+			while (chunkBytes > maxBuffer && chunks.length > 1) {
+				chunkBytes -= chunks.shift()!.length;
+			}
+
+			onChunk?.(text);
+		},
+		close() {
+			fullOutputStream?.end();
+		},
+	});
+
+	return Object.assign(sink, {
+		dump(annotation?: string) {
+			if (annotation) {
+				chunks.push(`\n\n${annotation}`);
+			}
+			const full = chunks.join("");
+			const { content, truncated } = truncateTail(full);
+			return { output: truncated ? content : full, truncated, fullOutputPath: fullOutputPath };
+		},
+	});
+}
+
+/**
+ * Execute a bash command with optional streaming and cancellation support.
+ *
+ * Features:
+ * - Streams sanitized output via onChunk callback
+ * - Writes large output to temp file for later retrieval
+ * - Supports cancellation via AbortSignal
+ * - Sanitizes output (strips ANSI, removes binary garbage, normalizes newlines)
+ * - Truncates output if it exceeds the default max bytes
+ *
+ * @param command - The bash command to execute
+ * @param options - Optional streaming callback and abort signal
+ * @returns Promise resolving to execution result
+ */
+export async function executeBash(command: string, options?: BashExecutorOptions): Promise<BashResult> {
+	const { shell, args, env, prefix } = getShellConfig();
+
+	// Get or create shell snapshot (for aliases, functions, options)
+	const snapshotPath = await getOrCreateSnapshot(shell, env);
+	const snapshotPrefix = getSnapshotSourceCommand(snapshotPath);
+
+	// Build final command: snapshot + prefix + command
+	const prefixedCommand = prefix ? `${prefix} ${command}` : command;
+	const finalCommand = `${snapshotPrefix}${prefixedCommand}`;
+
+	using signal = new ScopeSignal(options);
+
+	const child: Subprocess = Bun.spawn([shell, ...args, finalCommand], {
+		cwd: options?.cwd,
+		stdin: "ignore",
+		stdout: "pipe",
+		stderr: "pipe",
+		env,
+	});
+
+	signal.catch(() => {
+		killProcessTree(child.pid);
+	});
+
+	const sink = createOutputSink(DEFAULT_MAX_BYTES, DEFAULT_MAX_BYTES * 2, options?.onChunk);
+
+	const writer = sink.getWriter();
+	try {
+		async function pumpStream(readable: ReadableStream<Uint8Array>) {
+			const reader = readable.pipeThrough(createSanitizer()).getReader();
+			try {
+				while (true) {
+					const { done, value } = await reader.read();
+					if (done) break;
+					await writer.write(value);
+				}
+			} finally {
+				reader.releaseLock();
+			}
+		}
+		await Promise.all([
+			pumpStream(child.stdout as ReadableStream<Uint8Array>),
+			pumpStream(child.stderr as ReadableStream<Uint8Array>),
+		]);
+	} finally {
+		await writer.close();
+	}
+
+	// Non-zero exit codes or signal-killed processes are considered cancelled if killed via signal
+	const exitCode = await child.exited;
+
+	const cancelled = exitCode === null || (exitCode !== 0 && (options?.signal?.aborted ?? false));
+
+	if (signal.timedOut()) {
+		const secs = Math.round(options!.timeout! / 1000);
+		return {
+			exitCode: undefined,
+			cancelled: true,
+			...sink.dump(`Command timed out after ${secs} seconds`),
+		};
+	}
+
+	return {
+		exitCode: cancelled ? undefined : exitCode,
+		cancelled,
+		...sink.dump(),
+	};
+}