my-pi/pi-web-search/auth.ts

/**
 * Search provider API keys live in Pi's own auth.json — one `api_key`
 * credential per provider id (`tavily`, `exa`, `keenable`) next to Pi's
 * model credentials. Key values follow upstream resolve-config-value
 * semantics (replicated here because the upstream resolver is not exported):
 *
 * - `!command` executes the rest of the string and uses trimmed stdout,
 *   cached for the process lifetime (same as upstream's command cache)
 * - `$NAME` / `${NAME}` interpolate environment variables; a credential-level
 *   `env` map wins over process env; a missing variable resolves to undefined
 * - `$$` and `$!` escape to literal `$` and `!`
 * - anything else is a literal key
 */

import { execSync } from "node:child_process";
import { readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";

export type SearchAuthProvider = "tavily" | "exa" | "keenable";
export const SEARCH_AUTH_PROVIDERS: readonly SearchAuthProvider[] = ["tavily", "exa", "keenable"];

/** Resolves a provider's API key; injectable so tests never touch real auth.json. */
export type SearchKeyResolver = (provider: SearchAuthProvider) => string | undefined;

export type SearchAuthEnvironment = Record<string, string | undefined>;

export interface SearchAuthOptions {
	/** auth.json location; defaults to `<agent dir>/auth.json`. */
	authPath?: string;
	/**
	 * Environment used for the agent dir and `$VAR` interpolation. When
	 * provided it fully replaces `process.env` so tests stay isolated.
	 */
	env?: SearchAuthEnvironment;
	/** Executes `!command` key references; defaults to a cached execSync runner. */
	runCommand?: (command: string) => string | undefined;
}

export interface SearchApiKeyCredential {
	type: "api_key";
	key?: string;
	env?: Record<string, string>;
}

const COMMAND_TIMEOUT_MS = 10_000;
const ENV_VAR_NAME = /^[A-Za-z_][A-Za-z0-9_]*$/;
const ENV_VAR_NAME_PREFIX = /^[A-Za-z_][A-Za-z0-9_]*/;
const commandResultCache = new Map<string, string | undefined>();

/** auth.json lives in Pi's agent dir: $PI_CODING_AGENT_DIR or ~/.pi/agent. */
export function getSearchAuthPath(env: SearchAuthEnvironment = process.env): string {
	const agentDir = env.PI_CODING_AGENT_DIR?.trim() || join(env.HOME?.trim() || homedir(), ".pi", "agent");
	return join(agentDir, "auth.json");
}

/** Read the raw `api_key` credential for a provider; undefined when absent or malformed. */
export function readSearchCredential(
	provider: SearchAuthProvider,
	options: SearchAuthOptions = {},
): SearchApiKeyCredential | undefined {
	let data: unknown;
	try {
		data = JSON.parse(readFileSync(options.authPath ?? getSearchAuthPath(options.env), "utf8"));
	} catch {
		return undefined;
	}
	if (!data || typeof data !== "object" || Array.isArray(data)) return undefined;
	const entry = (data as Record<string, unknown>)[provider];
	if (!entry || typeof entry !== "object" || Array.isArray(entry)) return undefined;
	const credential = entry as Record<string, unknown>;
	if (credential.type !== "api_key") return undefined;
	if (credential.key !== undefined && typeof credential.key !== "string") return undefined;
	const env = credential.env;
	const scopedEnv =
		env && typeof env === "object" && !Array.isArray(env)
			? Object.fromEntries(
					Object.entries(env as Record<string, unknown>).filter(
						(pair): pair is [string, string] => typeof pair[1] === "string",
					),
				)
			: undefined;
	return { type: "api_key", key: credential.key as string | undefined, env: scopedEnv };
}

type TemplatePart = { type: "literal"; value: string } | { type: "env"; name: string };

/** Mirrors upstream parseConfigValueTemplate: `$$`/`$!` escape, `${A}`/`$A` interpolate. */
export function parseAuthKeyTemplate(config: string): TemplatePart[] {
	const parts: TemplatePart[] = [];
	const appendLiteral = (value: string) => {
		if (!value) return;
		const previous = parts[parts.length - 1];
		if (previous?.type === "literal") {
			previous.value += value;
			return;
		}
		parts.push({ type: "literal", value });
	};
	let index = 0;
	while (index < config.length) {
		const dollarIndex = config.indexOf("$", index);
		if (dollarIndex < 0) {
			appendLiteral(config.slice(index));
			break;
		}
		appendLiteral(config.slice(index, dollarIndex));
		const nextChar = config[dollarIndex + 1];
		if (nextChar === "$" || nextChar === "!") {
			appendLiteral(nextChar);
			index = dollarIndex + 2;
			continue;
		}
		if (nextChar === "{") {
			const endIndex = config.indexOf("}", dollarIndex + 2);
			if (endIndex < 0) {
				appendLiteral("$");
				index = dollarIndex + 1;
				continue;
			}
			const name = config.slice(dollarIndex + 2, endIndex);
			if (ENV_VAR_NAME.test(name)) {
				parts.push({ type: "env", name });
			} else {
				appendLiteral(config.slice(dollarIndex, endIndex + 1));
			}
			index = endIndex + 1;
			continue;
		}
		const match = config.slice(dollarIndex + 1).match(ENV_VAR_NAME_PREFIX);
		if (match) {
			parts.push({ type: "env", name: match[0] });
			index = dollarIndex + 1 + match[0].length;
			continue;
		}
		appendLiteral("$");
		index = dollarIndex + 1;
	}
	return parts;
}

function defaultRunCommand(command: string): string | undefined {
	if (commandResultCache.has(command)) return commandResultCache.get(command);
	let value: string | undefined;
	try {
		const output = execSync(command, {
			encoding: "utf-8",
			timeout: COMMAND_TIMEOUT_MS,
			stdio: ["ignore", "pipe", "ignore"],
		});
		value = output.trim() || undefined;
	} catch {
		value = undefined;
	}
	commandResultCache.set(command, value);
	return value;
}

/** Mirrors upstream resolveConfigValue, including credential-scoped env precedence. */
export function resolveAuthKeyValue(
	config: string,
	credentialEnv: Record<string, string> | undefined,
	options: SearchAuthOptions = {},
): string | undefined {
	if (config.startsWith("!")) {
		return (options.runCommand ?? defaultRunCommand)(config.slice(1));
	}
	const ambient = options.env ?? process.env;
	let resolved = "";
	for (const part of parseAuthKeyTemplate(config)) {
		if (part.type === "literal") {
			resolved += part.value;
			continue;
		}
		const value = credentialEnv?.[part.name] || ambient[part.name];
		if (value === undefined) return undefined;
		resolved += value;
	}
	return resolved;
}

/** Provider API key from auth.json, or undefined when not configured. */
export function getSearchApiKey(
	provider: SearchAuthProvider,
	options: SearchAuthOptions = {},
): string | undefined {
	const credential = readSearchCredential(provider, options);
	if (!credential?.key) return undefined;
	const resolved = resolveAuthKeyValue(credential.key, credential.env, options);
	return resolved?.trim() || undefined;
}

/** A credential-scoped config value (e.g. KEENABLE_API_URL), or undefined. */
export function getSearchCredentialEnv(
	provider: SearchAuthProvider,
	name: string,
	options: SearchAuthOptions = {},
): string | undefined {
	const value = readSearchCredential(provider, options)?.env?.[name]?.trim();
	return value || undefined;
}