import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.js";
import { discoverAndLoadExtensions } from "../src/core/extensions/loader.js";
import { ExtensionRunner } from "../src/core/extensions/runner.js";
import { ModelRegistry } from "../src/core/model-registry.js";
import { SessionManager } from "../src/core/session-manager.js";

describe("Before Agent Message Send Event", () => {
	let tempDir: string;
	let extensionsDir: string;

	beforeEach(() => {
		tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-agent-msg-ext-test-"));
		extensionsDir = path.join(tempDir, "extensions");
		fs.mkdirSync(extensionsDir);
	});

	afterEach(() => fs.rmSync(tempDir, { recursive: true, force: true }));

	async function createRunner(...extensions: string[]) {
		fs.rmSync(extensionsDir, { recursive: true, force: true });
		fs.mkdirSync(extensionsDir);
		for (let i = 0; i < extensions.length; i++) {
			fs.writeFileSync(path.join(extensionsDir, `e${i}.ts`), extensions[i]!);
		}
		const result = await discoverAndLoadExtensions([], tempDir, tempDir);
		const sm = SessionManager.inMemory();
		const mr = ModelRegistry.create(AuthStorage.create(path.join(tempDir, "auth.json")));
		return new ExtensionRunner(result.extensions, result.runtime, tempDir, sm, mr);
	}

	it("returns undefined when no handlers registered", async () => {
		const runner = await createRunner();
		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "parent-session-123",
			message: "hello parent",
			receiverRole: "parent",
		});
		expect(result).toBeUndefined();
	});

	it("allows an extension to append context usage when sending to parent", async () => {
		const ext = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event, ctx) => {
		if (event.receiverRole === "parent") {
			const usage = ctx.getContextUsage();
			const tokenText = usage ? \`[\${usage.tokens}/\${usage.contextWindow}]\` : "[no-usage]";
			return { message: \`\${event.message} \${tokenText}\` };
		}
	});
}
`;
		const runner = await createRunner(ext);
		runner.bindCore(
			{} as any,
			{
				getContextUsage: () => ({ tokens: 5000, contextWindow: 128000, percent: 3.90625 }),
			} as any,
		);

		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "parent-1",
			message: "work finished",
			receiverRole: "parent",
		});
		expect(result).toEqual({ message: "work finished [5000/128000]" });
	});

	it("does not modify message when receiverRole is not parent", async () => {
		const ext = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event, ctx) => {
		if (event.receiverRole === "parent") {
			return { message: event.message + " (to-parent)" };
		}
	});
}
`;
		const runner = await createRunner(ext);
		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "sibling-1",
			message: "hey sibling",
			receiverRole: "sibling",
			receiverName: "sibling-1",
		});
		expect(result).toBeUndefined();
	});

	it("chains modifications across multiple extensions", async () => {
		const ext1 = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event) => {
		return { message: event.message + " [first]" };
	});
}
`;
		const ext2 = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event) => {
		return { message: event.message + " [second]" };
	});
}
`;
		const runner = await createRunner(ext1, ext2);
		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "target-1",
			message: "base",
		});
		expect(result).toEqual({ message: "base [first] [second]" });
	});

	it("cancels message delivery when handler returns cancel: true", async () => {
		const ext = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event) => {
		if (event.message.includes("secret")) {
			return { cancel: true, reason: "secrets cannot be sent" };
		}
	});
}
`;
		const runner = await createRunner(ext);
		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "parent-1",
			message: "contains secret data",
		});
		expect(result).toEqual({ cancel: true, reason: "secrets cannot be sent" });
	});

	it("handles extension errors gracefully without crashing", async () => {
		const failingExt = `
export default function(pi) {
	pi.on("before_agent_message_send", async () => {
		throw new Error("handler exploded");
	});
}
`;
		const okExt = `
export default function(pi) {
	pi.on("before_agent_message_send", async (event) => {
		return { message: event.message + " ok" };
	});
}
`;
		const runner = await createRunner(failingExt, okExt);
		const errors: any[] = [];
		runner.onError((e) => errors.push(e));

		const result = await runner.emitBeforeAgentMessageSend({
			type: "before_agent_message_send",
			target: "parent-1",
			message: "test",
		});
		expect(result).toEqual({ message: "test ok" });
		expect(errors.length).toBe(1);
		expect(errors[0].event).toBe("before_agent_message_send");
	});
});

import { createHarness, type Harness } from "./suite/harness.js";

describe("AgentSession agent message host request with extensions and context usage", () => {
	let harness: Harness | undefined;

	afterEach(() => {
		if (harness) {
			harness.cleanup();
			harness = undefined;
		}
	});

	it("intercepts agent_message.send and appends context usage via extension", async () => {
		let sentMessage = "";
		let sentRole: any;
		const mockController = {
			listAgents: () => ({ parent: null, siblings: [], children: [] }),
			roster: () => ({
				entries: [
					{
						id: "parent-id",
						name: "parent",
						relationship: "parent" as const,
						depth: 0,
						status: "running" as const,
					},
				],
			}),
			sendAgentMessage: async (input: any) => {
				sentMessage = input.message;
				sentRole = input.receiverRole;
				return {
					id: "msg-1",
					source: "agent_message" as const,
					target: { activeSessionId: input.target, sessionId: "sess-1" },
					message: input.message,
					deliveryStatus: "delivered" as const,
				};
			},
		};

		const extFactory = (pi: any) => {
			pi.on("before_agent_message_send", async (event: any, ctx: any) => {
				if (event.receiverRole === "parent") {
					const usage = ctx.getContextUsage();
					const tokenPart = usage ? `[tokens: ${usage.tokens}/${usage.contextWindow}]` : "[no-usage]";
					return { message: `${event.message} ${tokenPart}` };
				}
			});
		};

		harness = await createHarness({
			agentMessageController: mockController as any,
			extensionFactories: [extFactory],
		});

		await harness.session.handleAgentMessageHostRequest("agent_message.send", {
			target: "parent-id",
			message: "task complete",
			receiverRole: "parent",
		});

		expect(sentMessage).toContain("task complete [tokens:");
		expect(sentRole).toBe("parent");
	});

	it("throws when extension cancels agent_message.send", async () => {
		const mockController = {
			listAgents: () => ({ parent: null, siblings: [], children: [] }),
			sendAgentMessage: async () => {
				throw new Error("should not be called");
			},
		};

		const extFactory = (pi: any) => {
			pi.on("before_agent_message_send", async () => {
				return { cancel: true, reason: "blocked by policy" };
			});
		};

		harness = await createHarness({
			agentMessageController: mockController as any,
			extensionFactories: [extFactory],
		});

		await expect(
			harness.session.handleAgentMessageHostRequest("agent_message.send", {
				target: "parent-id",
				message: "forbidden text",
				receiverRole: "parent",
			}),
		).rejects.toThrow("blocked by policy");
	});

	it("exposes rlm.context_usage in kernel host handlers", async () => {
		harness = await createHarness();
		const handlers = (harness.session as any)._createKernelHostHandlers();
		expect(handlers["rlm.context_usage"]).toBeDefined();
		expect(handlers["rlm.get_context_usage"]).toBeDefined();

		const res = await handlers["rlm.context_usage"]();
		expect(res).toBeDefined();
		expect(res.usage).toBeDefined();
		expect(res.usage.context_window).toBeGreaterThan(0);
	});
});