my-pi/tests/search-auth.test.ts

import assert from "node:assert/strict";
import { existsSync } from "node:fs";
import { mkdtemp, readFile, stat, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { spawnSync } from "node:child_process";
import test from "node:test";

import {
	getSearchApiKey,
	getSearchAuthPath,
	getSearchCredentialEnv,
	parseAuthKeyTemplate,
	resolveAuthKeyValue,
} from "../pi-web-search/auth.ts";

async function withAuthFile(
	data: Record<string, unknown>,
	fn: (authPath: string) => Promise<void> | void,
): Promise<void> {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-"));
	const authPath = join(dir, "auth.json");
	await writeFile(authPath, JSON.stringify(data));
	await fn(authPath);
}

test("auth path follows PI_CODING_AGENT_DIR and falls back to ~/.pi/agent", () => {
	assert.equal(
		getSearchAuthPath({ PI_CODING_AGENT_DIR: "/tmp/agent" }),
		"/tmp/agent/auth.json",
	);
	assert.equal(getSearchAuthPath({ HOME: "/tmp/home" }), "/tmp/home/.pi/agent/auth.json");
});

test("literal api_key credentials resolve verbatim", async () => {
	await withAuthFile(
		{
			tavily: { type: "api_key", key: "tvly-test" },
			exa: { type: "api_key", key: "  exa-test  " },
			keenable: { type: "api_key", key: "" },
			anthropic: { type: "oauth", access: "a", refresh: "r", expires: 1 },
		},
		(authPath) => {
			assert.equal(getSearchApiKey("tavily", { authPath }), "tvly-test");
			assert.equal(getSearchApiKey("exa", { authPath }), "exa-test");
			assert.equal(getSearchApiKey("keenable", { authPath }), undefined);
		},
	);
});

test("missing file, malformed json, non-api_key entries all resolve to undefined", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-"));
	assert.equal(getSearchApiKey("tavily", { authPath: join(dir, "absent.json") }), undefined);
	const bad = join(dir, "bad.json");
	await writeFile(bad, "not json{");
	assert.equal(getSearchApiKey("tavily", { authPath: bad }), undefined);
	await withAuthFile({ tavily: { type: "oauth", access: "a", refresh: "r", expires: 1 } }, (authPath) => {
		assert.equal(getSearchApiKey("tavily", { authPath }), undefined);
	});
	await withAuthFile({ tavily: "literal-string-entry" }, (authPath) => {
		assert.equal(getSearchApiKey("tavily", { authPath }), undefined);
	});
});

test("$VAR and ${VAR} interpolate; credential env beats ambient env", async () => {
	await withAuthFile(
		{
			tavily: { type: "api_key", key: "$TAVILY_KEY" },
			exa: { type: "api_key", key: "prefix-${EXA_KEY}-suffix" },
			keenable: { type: "api_key", key: "$SHARED", env: { SHARED: "scoped-wins" } },
		},
		(authPath) => {
			const env = { TAVILY_KEY: "from-env", EXA_KEY: "braced", SHARED: "ambient-loses" };
			assert.equal(getSearchApiKey("tavily", { authPath, env }), "from-env");
			assert.equal(getSearchApiKey("exa", { authPath, env }), "prefix-braced-suffix");
			assert.equal(getSearchApiKey("keenable", { authPath, env }), "scoped-wins");
		},
	);
});

test("a missing env variable resolves the whole key to undefined", async () => {
	await withAuthFile({ tavily: { type: "api_key", key: "$NOPE_MISSING" } }, (authPath) => {
		assert.equal(getSearchApiKey("tavily", { authPath, env: {} }), undefined);
	});
});

test("$$ and $! escape to literal characters", async () => {
	await withAuthFile(
		{ tavily: { type: "api_key", key: "cost-is-$$5-and-$!literal" } },
		(authPath) => {
			assert.equal(getSearchApiKey("tavily", { authPath, env: {} }), "cost-is-$5-and-!literal");
		},
	);
});

test("!command executes and caches trimmed stdout", async () => {
	const authPath = join(await mkdtemp(join(tmpdir(), "my-pi-auth-")), "auth.json");
	await writeFile(
		authPath,
		JSON.stringify({ tavily: { type: "api_key", key: "!printf secret-from-cmd" } }),
	);
	assert.equal(getSearchApiKey("tavily", { authPath }), "secret-from-cmd");

	let calls = 0;
	const runCommand = (command: string) => {
		calls += 1;
		assert.equal(command, "get-key --flag");
		return "injected";
	};
	await withAuthFile({ exa: { type: "api_key", key: "!get-key --flag" } }, (otherPath) => {
		assert.equal(getSearchApiKey("exa", { authPath: otherPath, runCommand }), "injected");
		assert.equal(getSearchApiKey("exa", { authPath: otherPath, runCommand }), "injected");
	});
});

test("a failing !command resolves to undefined", async () => {
	await withAuthFile({ tavily: { type: "api_key", key: "!exit 3" } }, (authPath) => {
		assert.equal(getSearchApiKey("tavily", { authPath }), undefined);
	});
});

test("credential env map exposes non-key config like KEENABLE_API_URL", async () => {
	await withAuthFile(
		{ keenable: { type: "api_key", key: "k", env: { KEENABLE_API_URL: "https://keenable.example" } } },
		(authPath) => {
			assert.equal(
				getSearchCredentialEnv("keenable", "KEENABLE_API_URL", { authPath }),
				"https://keenable.example",
			);
			assert.equal(getSearchCredentialEnv("keenable", "MISSING", { authPath }), undefined);
		},
	);
});

test("template parser handles env names, braces and escapes", () => {
	assert.deepEqual(parseAuthKeyTemplate("$A_${B}-$$-$!-x"), [
		{ type: "env", name: "A_" },
		{ type: "env", name: "B" },
		{ type: "literal", value: "-$-!-x" },
	]);
});

const AUTH_MJS = new URL("../scripts/search-auth.mjs", import.meta.url).pathname;
const CONFIGURE_SH = new URL("../configure.sh", import.meta.url).pathname;

function runMjs(args: string[], env: Record<string, string> = {}, input?: string) {
	return spawnSync("node", [AUTH_MJS, ...args], {
		encoding: "utf8",
		env: { ...process.env, ...env },
		...(input === undefined ? {} : { input }),
	});
}

test("search-auth.mjs --set-stdin writes api_key entries and preserves other providers", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(
		authPath,
		JSON.stringify({ anthropic: { type: "oauth", access: "tok", refresh: "ref", expires: 42, accountId: "acc" } }),
	);
	const first = runMjs(["--auth", authPath, "--set-stdin", "tavily"], {}, "tvly-new\n");
	assert.equal(first.status, 0, first.stderr);
	const second = runMjs(["--auth", authPath, "--set-stdin", "exa"], {}, "  !op read x  \n");
	assert.equal(second.status, 0, second.stderr);
	const written = JSON.parse(await readFile(authPath, "utf8")) as Record<string, any>;
	assert.deepEqual(written.anthropic, { type: "oauth", access: "tok", refresh: "ref", expires: 42, accountId: "acc" });
	assert.deepEqual(written.tavily, { type: "api_key", key: "tvly-new" });
	assert.deepEqual(written.exa, { type: "api_key", key: "!op read x" });
	assert.equal((await stat(authPath)).mode & 0o777, 0o600);
	assert.equal(getSearchApiKey("tavily", { authPath }), "tvly-new");
});

test("search-auth.mjs --set is rejected: keys never travel in argv", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	const result = runMjs(["--auth", authPath, "--set", "tavily=tvly-leak"]);
	assert.notEqual(result.status, 0);
	assert.match(result.stderr, /命令行参数|--set-stdin/);
	assert.equal(existsSync(authPath), false);
});

test("search-auth.mjs allows only one --set-stdin and rejects empty stdin", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	const dup = runMjs(["--auth", authPath, "--set-stdin", "tavily", "--set-stdin", "exa"], {}, "key");
	assert.notEqual(dup.status, 0);
	const empty = runMjs(["--auth", authPath, "--set-stdin", "tavily"], {}, "  \n");
	assert.notEqual(empty.status, 0);
	assert.equal(existsSync(authPath), false);
});

test("search-auth.mjs preserves an existing credential env map on --set-stdin", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(
		authPath,
		JSON.stringify({ keenable: { type: "api_key", key: "old", env: { KEENABLE_API_URL: "https://k.example" } } }),
	);
	const result = runMjs(["--auth", authPath, "--set-stdin", "keenable"], {}, "new-key");
	assert.equal(result.status, 0, result.stderr);
	const written = JSON.parse(await readFile(authPath, "utf8")) as Record<string, any>;
	assert.deepEqual(written.keenable, {
		type: "api_key",
		key: "new-key",
		env: { KEENABLE_API_URL: "https://k.example" },
	});
});

test("search-auth.mjs creates a missing auth.json and honors PI_CODING_AGENT_DIR", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const agentDir = join(dir, "agent");
	const result = runMjs(["--set-stdin", "keenable"], { PI_CODING_AGENT_DIR: agentDir }, "keen-1");
	assert.equal(result.status, 0, result.stderr);
	const written = JSON.parse(await readFile(join(agentDir, "auth.json"), "utf8")) as Record<string, any>;
	assert.deepEqual(written.keenable, { type: "api_key", key: "keen-1" });
});

test("search-auth.mjs --unset removes only the named provider", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(authPath, JSON.stringify({ tavily: { type: "api_key", key: "t" }, exa: { type: "api_key", key: "e" } }));
	const result = runMjs(["--auth", authPath, "--unset", "tavily"]);
	assert.equal(result.status, 0, result.stderr);
	const written = JSON.parse(await readFile(authPath, "utf8")) as Record<string, any>;
	assert.equal(written.tavily, undefined);
	assert.deepEqual(written.exa, { type: "api_key", key: "e" });
});

test("search-auth.mjs refuses to clobber a malformed auth.json", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(authPath, "{oops");
	const result = runMjs(["--auth", authPath, "--set-stdin", "tavily"], {}, "x");
	assert.notEqual(result.status, 0);
	assert.match(result.stderr, /不是合法 JSON|abort|clobber/i);
	assert.equal(await readFile(authPath, "utf8"), "{oops");
});

test("search-auth.mjs --migrate imports search.env keys without overwriting existing entries", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(authPath, JSON.stringify({ tavily: { type: "api_key", key: "existing" } }));
	const envPath = join(dir, "search.env");
	await writeFile(envPath, "export TAVILY_API_KEY=old\nEXA_API_KEY=exa-file\nKEENABLE_API_KEY=bad$(cmd)\n");
	const result = runMjs(["--auth", authPath, "--migrate", envPath]);
	assert.equal(result.status, 0, result.stderr);
	assert.match(result.stdout, /已迁移 exa/);
	assert.match(result.stdout, /tavily/);
	const written = JSON.parse(await readFile(authPath, "utf8")) as Record<string, any>;
	assert.equal(written.tavily.key, "existing");
	assert.equal(written.exa.key, "exa-file");
	assert.equal(written.keenable, undefined);
});

test("search-auth.mjs --has and --list never print key material", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	await writeFile(authPath, JSON.stringify({ tavily: { type: "api_key", key: "supersecret" } }));
	assert.equal(runMjs(["--auth", authPath, "--has", "tavily"]).status, 0);
	assert.equal(runMjs(["--auth", authPath, "--has", "exa"]).status, 1);
	const list = runMjs(["--auth", authPath, "--list"]);
	assert.equal(list.status, 0);
	assert.match(list.stdout, /tavily: 已配置/);
	assert.doesNotMatch(list.stdout, /supersecret/);
});

test("configure.sh search writes keys through search-auth.mjs via stdin", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	const first = spawnSync("sh", [CONFIGURE_SH, "search", "--auth", authPath, "--tavily", "-"], {
		encoding: "utf8",
		input: "tvly-sh\n",
	});
	assert.equal(first.status, 0, first.stderr);
	const second = spawnSync("sh", [CONFIGURE_SH, "search", "--auth", authPath, "--exa", "-"], {
		encoding: "utf8",
		input: "exa-sh",
	});
	assert.equal(second.status, 0, second.stderr);
	const written = JSON.parse(await readFile(authPath, "utf8")) as Record<string, any>;
	assert.deepEqual(written.tavily, { type: "api_key", key: "tvly-sh" });
	assert.deepEqual(written.exa, { type: "api_key", key: "exa-sh" });
	assert.equal((await stat(authPath)).mode & 0o777, 0o600);
});

test("configure.sh search rejects a plaintext key argument", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	const result = spawnSync("sh", [CONFIGURE_SH, "search", "--auth", authPath, "--tavily", "tvly-leak"], {
		encoding: "utf8",
	});
	assert.notEqual(result.status, 0);
	assert.match(result.stderr, /命令行参数/);
	assert.equal(existsSync(authPath), false);
});

test("configure.sh search allows at most one stdin provider per call", async () => {
	const dir = await mkdtemp(join(tmpdir(), "my-pi-auth-write-"));
	const authPath = join(dir, "auth.json");
	const result = spawnSync("sh", [CONFIGURE_SH, "search", "--auth", authPath, "--tavily", "-", "--exa", "-"], {
		encoding: "utf8",
		input: "one-key",
	});
	assert.notEqual(result.status, 0);
	assert.match(result.stderr, /最多一个/);
});