pi-ext

import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
	DefaultResourceLoader,
	SessionManager,
} from "@earendil-works/pi-coding-agent";
import { describe, expect, it, vi } from "vitest";
import {
	type HeadlessSession,
	makeResolveModel,
	makeSessionFactory,
	makeSubagentSessionManager,
} from "../drivers.ts";
import type { ResolvedDynamicExtension } from "../dynamic-types.ts";

// These tests assert the OPTION CONSTRUCTION + request/response MAPPING of the
// real drivers, including resource isolation, initial tool sets, and persistent
// parent-linked session management, by
// injecting fake SDK constructors/functions and
// inspecting how they're called. Real resource loaders and session managers
// additionally prove applicable instruction loading and child-history isolation
// without launching a model.

function fakeSession(): HeadlessSession {
	let active: string[] = [];
	return {
		subscribe: () => () => {},
		prompt: async () => {},
		steer: () => {},
		abort: () => {},
		dispose: () => {},
		setActiveToolsByName: vi.fn((names: string[]) => {
			active = names;
		}),
		getActiveToolNames: () => active,
		bindExtensions: vi.fn().mockResolvedValue(undefined),
	};
}

describe("makeSessionFactory", () => {
	it("constructs an isolated, reloaded session with the goal-proven option set", async () => {
		const reload = vi.fn().mockResolvedValue(undefined);
		const loaderOpts: unknown[] = [];
		class FakeLoader {
			reload = reload;
			constructor(opts: unknown) {
				loaderOpts.push(opts);
			}
		}
		// The real createAgentSession resolves to `{ session }` (destructured).
		const child = fakeSession();
		const createAgentSession = vi.fn(async () => ({ session: child }));

		const factory = makeSessionFactory({
			createAgentSession,
			DefaultResourceLoader: FakeLoader as any,
			getAgentDir: () => "/agent",
			cwd: "/work",
		});

		const sessionManager = { id: "fresh-inmemory" };
		const customTools = [{ name: "structured_output" }];
		const session = await factory({
			model: { id: "m" },
			thinkingLevel: "low",
			tools: ["read", "structured_output"],
			customTools: customTools as any,
			sessionManager,
		});

		// Controlled resource discovery retains context files; reload precedes construction.
		expect(loaderOpts[0]).toEqual({
			cwd: "/work",
			agentDir: "/agent",
			noExtensions: true,
			additionalExtensionPaths: undefined,
			noSkills: true,
			noPromptTemplates: true,
			noThemes: true,
			noContextFiles: false,
		});
		expect(reload).toHaveBeenCalledTimes(1);
		expect(reload.mock.invocationCallOrder[0]).toBeLessThan(
			createAgentSession.mock.invocationCallOrder[0],
		);

		// The full option set, with every correction asserted.
		const opts = createAgentSession.mock.calls[0][0] as Record<string, unknown>;
		expect(opts.cwd).toBe("/work");
		expect(opts.model).toEqual({ id: "m" });
		expect(opts).not.toHaveProperty("modelRegistry");
		expect(opts.resourceLoader).toBeInstanceOf(FakeLoader);
		expect(opts.thinkingLevel).toBe("low");
		expect(opts.sessionManager).toBe(sessionManager); // manager identity preserved
		expect(opts).not.toHaveProperty("tools"); // extensions retain their full registry
		expect(opts.customTools).toBe(customTools);
		expect(child.setActiveToolsByName).toHaveBeenCalledWith([
			"read",
			"structured_output",
		]);
		expect(child.bindExtensions).toHaveBeenCalledWith({ mode: "print" });
		expect(
			vi.mocked(child.setActiveToolsByName).mock.invocationCallOrder[0],
		).toBeLessThan(vi.mocked(child.bindExtensions).mock.invocationCallOrder[0]);

		// Returns the destructured `.session`, not the wrapper.
		expect(session).toBe(child);
	});

	it("loads extensions and skills independently while preserving other isolation", async () => {
		const loaderOpts: Record<string, unknown>[] = [];
		class FakeLoader {
			reload = vi.fn().mockResolvedValue(undefined);
			constructor(opts: Record<string, unknown>) {
				loaderOpts.push(opts);
			}
		}
		const child = fakeSession();
		const createAgentSession = vi.fn(async () => ({ session: child }));

		const factory = makeSessionFactory({
			createAgentSession,
			DefaultResourceLoader: FakeLoader as any,
			getAgentDir: () => "/agent",
			cwd: "/work",
			loadExtensions: true,
			loadSkills: true,
		});
		await factory({
			model: {},
			tools: ["external_tool", "structured_output"],
			customTools: [],
			sessionManager: {},
		});

		expect(loaderOpts[0].noExtensions).toBe(false);
		expect(loaderOpts[0].noSkills).toBe(false);
		expect(child.setActiveToolsByName).toHaveBeenCalledWith([
			"external_tool",
			"structured_output",
			"read",
		]);
		expect(loaderOpts[0].noPromptTemplates).toBe(true);
		expect(loaderOpts[0].noThemes).toBe(true);
		expect(loaderOpts[0].noContextFiles).toBe(false);
	});

	it("appends ordered dynamic snapshots and builder guidance after normal extensions", async () => {
		const loaderOpts: Record<string, unknown>[] = [];
		class FakeLoader {
			reload = vi.fn().mockResolvedValue(undefined);
			constructor(opts: Record<string, unknown>) {
				loaderOpts.push(opts);
			}
			getExtensions() {
				return { extensions: [], runtime: {} };
			}
		}
		const factory = makeSessionFactory({
			createAgentSession: vi.fn(async () => ({ session: fakeSession() })),
			DefaultResourceLoader: FakeLoader,
			getAgentDir: () => "/agent",
			cwd: "/work",
			extensionPaths: ["/normal/index.ts"],
		});
		const dynamicExtensions: ResolvedDynamicExtension[] = [
			{
				name: "alpha-capability",
				description: "Handles alpha values.",
				revision: "a".repeat(64),
				selector: `alpha-capability@${"a".repeat(64)}`,
				canonicalPath: "/catalog/alpha-capability/revision",
				sourcePath: "/snapshot/alpha",
				loadPath: "/snapshot/alpha/index.ts",
				overrides: [],
			},
			{
				name: "beta-capability",
				description: "Handles beta values.",
				revision: "b".repeat(64),
				selector: `beta-capability@${"b".repeat(64)}`,
				canonicalPath: "/catalog/beta-capability/revision",
				sourcePath: "/snapshot/beta",
				loadPath: "/snapshot/beta/index.ts",
				overrides: [],
			},
		];

		await factory({
			model: {},
			tools: ["structured_output"],
			customTools: [],
			extensionPaths: dynamicExtensions.map((extension) => extension.loadPath),
			dynamicExtensions,
			appendSystemPrompt: ["builder contract"],
			sessionManager: {},
		});

		expect(loaderOpts[0].additionalExtensionPaths).toEqual([
			"/normal/index.ts",
			"/snapshot/alpha/index.ts",
			"/snapshot/beta/index.ts",
		]);
		expect(loaderOpts[0].appendSystemPrompt).toBeUndefined();
		const append = loaderOpts[0].appendSystemPromptOverride as (
			base: string[],
		) => string[];
		expect(append(["user append instructions"])).toEqual([
			"user append instructions",
			"builder contract",
		]);
	});

	it("blocks undeclared dynamic collisions and applies exact declared overrides", async () => {
		const normalTools = new Map([["shared", { normal: true }]]);
		const dynamicTools = new Map([["shared", { dynamic: true }]]);
		class FakeLoader {
			reload = vi.fn().mockResolvedValue(undefined);
			getExtensions() {
				return {
					extensions: [
						{ path: "/normal/index.ts", tools: normalTools },
						{ path: "/snapshot/index.ts", tools: dynamicTools },
					],
					runtime: {},
				};
			}
		}
		const extension: ResolvedDynamicExtension = {
			name: "shared-capability",
			description: "Handles shared values.",
			revision: "a".repeat(64),
			selector: `shared-capability@${"a".repeat(64)}`,
			canonicalPath: "/catalog/shared-capability/revision",
			sourcePath: "/snapshot",
			loadPath: "/snapshot/index.ts",
			overrides: [],
		};
		const factory = makeSessionFactory({
			createAgentSession: vi.fn(async () => ({ session: fakeSession() })),
			DefaultResourceLoader: FakeLoader,
			getAgentDir: () => "/agent",
			cwd: "/work",
		});
		const args = {
			model: {},
			tools: ["structured_output"],
			customTools: [],
			extensionPaths: [extension.loadPath],
			dynamicExtensions: [extension],
			sessionManager: {},
		};

		await expect(factory(args)).rejects.toThrow(/collides with tool:shared/u);
		extension.overrides = ["tool:shared"];
		await expect(factory(args)).resolves.toBeDefined();
		expect(normalTools.has("shared")).toBe(false);
		expect(dynamicTools.has("shared")).toBe(true);
	});

	it("requires builtin overrides and always rejects Ultra-reserved tool registrations", async () => {
		const dynamicTools = new Map([["read", { dynamic: true }]]);
		class FakeLoader {
			reload = vi.fn().mockResolvedValue(undefined);
			getExtensions() {
				return {
					extensions: [{ path: "/snapshot/index.ts", tools: dynamicTools }],
					runtime: {},
				};
			}
		}
		const extension: ResolvedDynamicExtension = {
			name: "builtin-override-capability",
			description: "Overrides one builtin deliberately.",
			revision: "a".repeat(64),
			selector: `builtin-override-capability@${"a".repeat(64)}`,
			canonicalPath: "/catalog/builtin-override-capability/revision",
			sourcePath: "/snapshot",
			loadPath: "/snapshot/index.ts",
			overrides: [],
		};
		const factory = makeSessionFactory({
			createAgentSession: vi.fn(async () => ({ session: fakeSession() })),
			DefaultResourceLoader: FakeLoader,
			getAgentDir: () => "/agent",
			cwd: "/work",
		});
		const args = {
			model: {},
			tools: ["structured_output"],
			customTools: [],
			extensionPaths: [extension.loadPath],
			dynamicExtensions: [extension],
			sessionManager: {},
		};

		await expect(factory(args)).rejects.toThrow(/collides with tool:read/u);
		extension.overrides = ["tool:read"];
		await expect(factory(args)).resolves.toBeDefined();
		dynamicTools.delete("read");
		dynamicTools.set("structured_output", { dynamic: true });
		await expect(factory(args)).rejects.toThrow(
			/cannot register reserved tool:structured_output/u,
		);
	});

	it("loads one static extension path list without broad discovery", async () => {
		const loaderOpts: Record<string, unknown>[] = [];
		class FakeLoader {
			reload = vi.fn().mockResolvedValue(undefined);
			constructor(opts: Record<string, unknown>) {
				loaderOpts.push(opts);
			}
		}
		const factory = makeSessionFactory({
			createAgentSession: vi.fn(async () => ({ session: fakeSession() })),
			DefaultResourceLoader: FakeLoader as any,
			getAgentDir: () => "/agent",
			cwd: "/work",
			extensionPaths: [
				"/extensions/alpha/index.ts",
				"/extensions/beta/index.ts",
			],
		});

		await factory({
			model: {},
			tools: ["external_tool", "structured_output"],
			customTools: [],
			sessionManager: {},
		});

		expect(loaderOpts[0]).toMatchObject({
			noExtensions: true,
			additionalExtensionPaths: [
				"/extensions/alpha/index.ts",
				"/extensions/beta/index.ts",
			],
		});
	});
});

describe("sub-agent instruction discovery", () => {
	it.each([false, true])(
		"loads real applicable instructions without replacing user system prompts (builder: %s)",
		async (builder) => {
			const root = await mkdtemp(join(tmpdir(), "ultra-instructions-"));
			const agentDir = join(root, "agent");
			const project = join(root, "project");
			const cwd = join(project, "nested");
			const sibling = join(project, "sibling");
			const applicable = [
				{
					path: join(agentDir, "AGENTS.md"),
					content: "Global user instructions.\n",
				},
				{
					path: join(project, "AGENTS.md"),
					content: "Project instructions.\n",
				},
				{
					path: join(cwd, "AGENTS.override.md"),
					content: "Nested override instructions.\n",
				},
			];
			try {
				await mkdir(agentDir, { recursive: true });
				await mkdir(cwd, { recursive: true });
				await mkdir(sibling, { recursive: true });
				for (const file of applicable) await writeFile(file.path, file.content);
				await writeFile(join(cwd, "AGENTS.md"), "Superseded instructions.\n");
				await writeFile(
					join(sibling, "AGENTS.md"),
					"Unrelated instructions.\n",
				);
				await writeFile(
					join(agentDir, "SYSTEM.md"),
					"User system instructions.",
				);
				await writeFile(
					join(agentDir, "APPEND_SYSTEM.md"),
					"User append instructions.",
				);
				const manager = SessionManager.inMemory(cwd);
				const createAgentSession = vi.fn(
					async (opts: Record<string, unknown>) => {
						const loader = opts.resourceLoader as DefaultResourceLoader;
						const files = loader.getAgentsFiles().agentsFiles;
						expect(files).toEqual(expect.arrayContaining(applicable));
						const paths = files.map((file) => file.path);
						expect(paths).not.toContain(join(cwd, "AGENTS.md"));
						expect(paths).not.toContain(join(sibling, "AGENTS.md"));
						expect(loader.getExtensions().extensions).toEqual([]);
						expect(loader.getSkills().skills).toEqual([]);
						expect(loader.getSystemPrompt()).toBe("User system instructions.");
						expect(loader.getAppendSystemPrompt()).toEqual([
							"User append instructions.",
							...(builder ? ["Builder contract."] : []),
						]);
						expect(opts.sessionManager).toBe(manager);
						expect(manager.buildSessionContext().messages).toEqual([]);
						return { session: fakeSession() };
					},
				);
				const factory = makeSessionFactory({
					createAgentSession,
					DefaultResourceLoader,
					getAgentDir: () => agentDir,
					cwd,
				});
				await factory({
					model: {},
					tools: ["structured_output"],
					customTools: [],
					sessionManager: manager,
					...(builder ? { appendSystemPrompt: ["Builder contract."] } : {}),
				});
				expect(createAgentSession).toHaveBeenCalledTimes(1);
			} finally {
				await rm(root, { recursive: true, force: true });
			}
		},
	);
});

describe("static sub-agent extension discovery", () => {
	it("loads exactly the selected extension paths", async () => {
		const root = await mkdtemp(join(tmpdir(), "ultra-extensions-"));
		const selected = [join(root, "selected-extension.ts")];
		try {
			await writeFile(
				selected[0],
				"export default function selectedExtension() {}\n",
			);
			const loader = new DefaultResourceLoader({
				cwd: root,
				agentDir: join(root, "agent"),
				noExtensions: true,
				additionalExtensionPaths: selected,
				noSkills: true,
				noPromptTemplates: true,
				noThemes: true,
				noContextFiles: true,
			});
			await loader.reload();

			expect(loader.getExtensions().errors).toEqual([]);
			expect(
				loader.getExtensions().extensions.map((extension) => extension.path),
			).toEqual(selected);
		} finally {
			await rm(root, { recursive: true, force: true });
		}
	});
});

describe("sub-agent skill discovery", () => {
	it("loads skills by default and supports complete opt-out", async () => {
		const root = await mkdtemp(join(tmpdir(), "ultra-skills-"));
		const agentDir = join(root, "agent");
		const skillDir = join(agentDir, "skills", "ultra-skill-probe");
		await mkdir(skillDir, { recursive: true });
		await writeFile(
			join(skillDir, "SKILL.md"),
			"---\nname: ultra-skill-probe\ndescription: Verifies Ultra skill discovery.\n---\n\n# Probe\n",
		);
		try {
			const loader = new DefaultResourceLoader({
				cwd: root,
				agentDir,
				noExtensions: true,
				noSkills: false,
				noPromptTemplates: true,
				noThemes: true,
				noContextFiles: true,
			});
			await loader.reload();
			expect(loader.getSkills().skills.map((skill) => skill.name)).toContain(
				"ultra-skill-probe",
			);

			const disabled = new DefaultResourceLoader({
				cwd: root,
				agentDir,
				noExtensions: true,
				noSkills: true,
				noPromptTemplates: true,
				noThemes: true,
				noContextFiles: true,
			});
			await disabled.reload();
			expect(disabled.getSkills().skills).toEqual([]);
		} finally {
			await rm(root, { recursive: true, force: true });
		}
	});
});

describe("makeSubagentSessionManager", () => {
	it("creates, names, parent-links, and reopens a persistent child", async () => {
		const root = await mkdtemp(join(tmpdir(), "ultra-sessions-"));
		const cwd = join(root, "work");
		const sessionDir = join(root, "sessions");
		await mkdir(cwd);
		try {
			const parent = SessionManager.create(cwd, sessionDir);
			parent.appendMessage({
				role: "user",
				content: "Parent conversation must not enter the child history.",
				timestamp: Date.now(),
			});
			const parentSession = parent.getSessionFile();
			expect(parentSession).toBeDefined();
			const manage = makeSubagentSessionManager({
				SessionManager,
				cwd,
				parentSession,
				retention: "all",
				sessionDir,
			});

			const fresh = manage({ name: "ultra review abc123 review#0" });
			const freshManager = fresh.manager as SessionManager;
			expect(fresh.resumed).toBe(false);
			expect(fresh.ref).toEqual({
				id: freshManager.getSessionId(),
				path: freshManager.getSessionFile(),
			});
			expect(freshManager.getHeader()?.parentSession).toBe(parentSession);
			expect(parent.buildSessionContext().messages).toHaveLength(1);
			expect(freshManager.buildSessionContext().messages).toEqual([]);
			expect(freshManager.getSessionName()).toBe(
				"ultra review abc123 review#0",
			);
			await writeFile(
				fresh.ref?.path as string,
				`${[freshManager.getHeader(), ...freshManager.getEntries()]
					.map((entry) => JSON.stringify(entry))
					.join("\n")}\n`,
			);

			const resumed = manage({ name: "ignored", retained: fresh.ref });
			expect(resumed.resumed).toBe(true);
			expect((resumed.manager as SessionManager).getSessionId()).toBe(
				freshManager.getSessionId(),
			);
		} finally {
			await rm(root, { recursive: true, force: true });
		}
	});

	it("replaces a retained child whose file disappeared", async () => {
		const root = await mkdtemp(join(tmpdir(), "ultra-sessions-"));
		const cwd = join(root, "work");
		const sessionDir = join(root, "sessions");
		await mkdir(cwd);
		try {
			const parent = SessionManager.create(cwd, sessionDir);
			const manage = makeSubagentSessionManager({
				SessionManager,
				cwd,
				parentSession: parent.getSessionFile(),
				retention: "all",
				sessionDir,
			});
			const missing = join(sessionDir, "missing.jsonl");
			const replacement = manage({
				name: "ultra replacement",
				retained: { id: "missing", path: missing },
			});

			expect(replacement.resumed).toBe(false);
			expect(replacement.ref?.id).not.toBe("missing");
			expect(replacement.ref?.path).not.toBe(missing);
		} finally {
			await rm(root, { recursive: true, force: true });
		}
	});

	it("stays stateless when opted out or when the parent has no session file", () => {
		for (const deps of [
			{ parentSession: "/parent.jsonl", retention: "none" as const },
			{ parentSession: undefined, retention: "all" as const },
		]) {
			const child = makeSubagentSessionManager({
				SessionManager,
				cwd: "/work",
				...deps,
			})({ name: "stateless" });
			expect(child.resumed).toBe(false);
			expect(child.ref).toBeUndefined();
			expect((child.manager as SessionManager).isPersisted()).toBe(false);
		}
	});
});

describe("makeResolveModel", () => {
	it("returns the default model when no ref is given", () => {
		const find = vi.fn();
		const resolve = makeResolveModel({
			modelRegistry: { find },
			defaultModel: { id: "host" },
		});

		expect(resolve(undefined)).toEqual({ id: "host" });
		expect(find).not.toHaveBeenCalled();
	});

	it("parses a provider/model ref and looks it up in the registry", () => {
		const found = { id: "MiniMax-M2" };
		const find = vi.fn(() => found);
		const resolve = makeResolveModel({
			modelRegistry: { find },
			defaultModel: { id: "host" },
		});

		expect(resolve("minimax/MiniMax-M2")).toBe(found);
		expect(find).toHaveBeenCalledWith("minimax", "MiniMax-M2");
	});

	it("throws on an unresolvable ref and on a malformed ref", () => {
		const resolve = makeResolveModel({
			modelRegistry: { find: () => undefined },
			defaultModel: {},
		});

		expect(() => resolve("minimax/nope")).toThrow(/not found/);
		expect(() => resolve("noslash")).toThrow(/provider\/model/);
		expect(() => resolve("/leading")).toThrow(/provider\/model/);
		expect(() => resolve("trailing/")).toThrow(/provider\/model/);
	});
});