oh-my-pi/packages/coding-agent/test/advisor-watchdog.test.ts

import { afterEach, describe, expect, it } from "bun:test";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk";
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
import { Snowflake } from "@oh-my-pi/pi-utils";

describe("advisor watchdog prompt discovery", () => {
	const tempDirs: string[] = [];

	afterEach(() => {
		for (const tempDir of tempDirs.splice(0)) {
			fs.rmSync(tempDir, { recursive: true, force: true });
		}
	});

	it("discovers and appends WATCHDOG.md to the advisor prompt", async () => {
		const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-advisor-watchdog-${Snowflake.next()}-`));
		tempDirs.push(tempDir);
		const cwd = path.join(tempDir, "project-root");
		fs.mkdirSync(cwd, { recursive: true });

		// Write a WATCHDOG.md file
		const watchdogContent = "Watchdog rule: Watch out for cheating on edits.";
		fs.writeFileSync(path.join(cwd, "WATCHDOG.md"), watchdogContent, "utf8");

		const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db"));
		authStorage.setRuntimeApiKey("openai", "test-key");
		const modelRegistry = new ModelRegistry(authStorage);

		const sessionManager = SessionManager.create(cwd, path.join(tempDir, "sessions"));
		const { session } = await createAgentSession({
			cwd,
			agentDir: tempDir,
			sessionManager,
			authStorage,
			modelRegistry,
			settings: (() => {
				const s = Settings.isolated({
					"async.enabled": false,
					"advisor.enabled": true,
				});
				s.setModelRole("advisor", "openai/gpt-4o-mini");
				return s;
			})(),
			model: getBundledModel("openai", "gpt-4o-mini"),
			disableExtensionDiscovery: true,
			skills: [],
			contextFiles: [],
			promptTemplates: [],
			slashCommands: [],
			enableMCP: false,
			enableLsp: false,
		});

		try {
			expect(session.isAdvisorActive()).toBe(true);
			const dump = session.formatAdvisorHistoryAsText();
			expect(dump).not.toBeNull();
			expect(dump).toContain("Especially pay attention to:");
			expect(dump).toContain("<attention>");
			expect(dump).toContain(watchdogContent);
			expect(dump).toContain("</attention>");
		} finally {
			await session.dispose();
		}
	});

	it("resolves nested folders and sorts by depth", async () => {
		const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-advisor-watchdog-${Snowflake.next()}-`));
		tempDirs.push(tempDir);
		const parentCwd = path.join(tempDir, "project-root");
		const childCwd = path.join(parentCwd, "subfolder");
		fs.mkdirSync(childCwd, { recursive: true });

		// Write two WATCHDOG.md files
		const parentWatchdogContent = "Parent watchdog rule.";
		const childWatchdogContent = "Child watchdog rule.";
		fs.writeFileSync(path.join(parentCwd, "WATCHDOG.md"), parentWatchdogContent, "utf8");
		fs.writeFileSync(path.join(childCwd, "WATCHDOG.md"), childWatchdogContent, "utf8");

		const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db"));
		authStorage.setRuntimeApiKey("openai", "test-key");
		const modelRegistry = new ModelRegistry(authStorage);

		const sessionManager = SessionManager.create(childCwd, path.join(tempDir, "sessions"));
		const { session } = await createAgentSession({
			cwd: childCwd,
			agentDir: tempDir,
			sessionManager,
			authStorage,
			modelRegistry,
			settings: (() => {
				const s = Settings.isolated({
					"async.enabled": false,
					"advisor.enabled": true,
				});
				s.setModelRole("advisor", "openai/gpt-4o-mini");
				return s;
			})(),
			model: getBundledModel("openai", "gpt-4o-mini"),
			disableExtensionDiscovery: true,
			skills: [],
			contextFiles: [],
			promptTemplates: [],
			slashCommands: [],
			enableMCP: false,
			enableLsp: false,
		});

		try {
			expect(session.isAdvisorActive()).toBe(true);
			const dump = session.formatAdvisorHistoryAsText();
			expect(dump).not.toBeNull();
			expect(dump).toContain("Especially pay attention to:");
			expect(dump).toContain("<attention>");
			expect(dump).toContain("</attention>");
			expect(dump).toContain(parentWatchdogContent);
			expect(dump).toContain(childWatchdogContent);
			// Check ordering: parent is farther (depth 1), child is closer (depth 0).
			// So parent watchdog should appear first, followed by child watchdog.
			const parentIndex = dump!.indexOf(parentWatchdogContent);
			const childIndex = dump!.indexOf(childWatchdogContent);
			expect(parentIndex).toBeGreaterThan(-1);
			expect(childIndex).toBeGreaterThan(-1);
			expect(parentIndex).toBeLessThan(childIndex);
		} finally {
			await session.dispose();
		}
	});

	it("discovers user-level and native project-level watchdog files", async () => {
		const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-advisor-watchdog-${Snowflake.next()}-`));
		tempDirs.push(tempDir);
		const cwd = path.join(tempDir, "project-root");
		const ompDir = path.join(cwd, ".omp");
		const userAgentDir = path.join(tempDir, "user-agent");
		fs.mkdirSync(cwd, { recursive: true });
		fs.mkdirSync(ompDir, { recursive: true });
		fs.mkdirSync(userAgentDir, { recursive: true });

		const userWatchdogContent = "User-level watchdog rule.";
		const nativeWatchdogContent = "Native project watchdog rule.";
		const standaloneWatchdogContent = "Standalone project watchdog rule.";

		fs.writeFileSync(path.join(userAgentDir, "WATCHDOG.md"), userWatchdogContent, "utf8");
		fs.writeFileSync(path.join(ompDir, "WATCHDOG.md"), nativeWatchdogContent, "utf8");
		fs.writeFileSync(path.join(cwd, "WATCHDOG.md"), standaloneWatchdogContent, "utf8");

		const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db"));
		authStorage.setRuntimeApiKey("openai", "test-key");
		const modelRegistry = new ModelRegistry(authStorage);

		const sessionManager = SessionManager.create(cwd, path.join(tempDir, "sessions"));
		const { session } = await createAgentSession({
			cwd,
			agentDir: userAgentDir,
			sessionManager,
			authStorage,
			modelRegistry,
			settings: (() => {
				const s = Settings.isolated({
					"async.enabled": false,
					"advisor.enabled": true,
				});
				s.setModelRole("advisor", "openai/gpt-4o-mini");
				return s;
			})(),
			model: getBundledModel("openai", "gpt-4o-mini"),
			disableExtensionDiscovery: true,
			skills: [],
			contextFiles: [],
			promptTemplates: [],
			slashCommands: [],
			enableMCP: false,
			enableLsp: false,
		});

		try {
			expect(session.isAdvisorActive()).toBe(true);
			const dump = session.formatAdvisorHistoryAsText();
			expect(dump).not.toBeNull();
			expect(dump).toContain(userWatchdogContent);
			expect(dump).toContain(nativeWatchdogContent);
			expect(dump).toContain(standaloneWatchdogContent);

			// Check ordering: user-level should appear first, then native project level (.omp/WATCHDOG.md has depth 0),
			// then standalone project level (cwd/WATCHDOG.md has depth 0).
			// Between native and standalone, they both have depth 0, so their relative order doesn't strictly matter
			// as long as user-level comes before both of them.
			const userIndex = dump!.indexOf(userWatchdogContent);
			const nativeIndex = dump!.indexOf(nativeWatchdogContent);
			const standaloneIndex = dump!.indexOf(standaloneWatchdogContent);

			expect(userIndex).toBeGreaterThan(-1);
			expect(nativeIndex).toBeGreaterThan(-1);
			expect(standaloneIndex).toBeGreaterThan(-1);

			expect(userIndex).toBeLessThan(nativeIndex);
			expect(userIndex).toBeLessThan(standaloneIndex);
		} finally {
			await session.dispose();
		}
	});
});