pi-ext

import { readFileSync, writeFileSync } from "node:fs";

const files = [
	"node_modules/@marcfargas/pi-test-harness/dist/session.js",
	"node_modules/@marcfargas/pi-test-harness/src/session.ts",
];

for (const file of files) {
	let text = readFileSync(file, "utf8");
	text = text.replace(
		'import { getModel } from "@earendil-works/pi-ai";',
		'import { getModel } from "@earendil-works/pi-ai/compat";',
	);
	text = text.replace(
		/import \{ createAgentSession, DefaultResourceLoader, SessionManager, SettingsManager,? \} from "@earendil-works\/pi-coding-agent";/,
		'import { createAgentSession, DefaultResourceLoader, ModelRuntime, SessionManager, SettingsManager } from "@earendil-works/pi-coding-agent";',
	);
	text = text.replace(
		/import \{\n\s*createAgentSession,\n\s*DefaultResourceLoader,\n\s*SessionManager,\n\s*SettingsManager,\n\s*type AgentSessionEvent,\n\} from "@earendil-works\/pi-coding-agent";/,
		'import {\n\tcreateAgentSession,\n\tDefaultResourceLoader,\n\tModelRuntime,\n\tSessionManager,\n\tSettingsManager,\n\ttype AgentSessionEvent,\n} from "@earendil-works/pi-coding-agent";',
	);
	text = text.replace(
		"\t// Create real session with in-memory persistence\n\tconst { session, extensionsResult } = await createAgentSession({",
		'\tconst modelRuntime = await ModelRuntime.create({ refreshOnCreate: false });\n\tawait modelRuntime.setRuntimeApiKey("openai", "test-key");\n\n\t// Create real session with in-memory persistence\n\tconst { session, extensionsResult } = await createAgentSession({',
	);
	text = text.replace(
		"    // Create real session with in-memory persistence\n    const { session, extensionsResult } = await createAgentSession({",
		'    const modelRuntime = await ModelRuntime.create({ refreshOnCreate: false });\n    await modelRuntime.setRuntimeApiKey("openai", "test-key");\n\n    // Create real session with in-memory persistence\n    const { session, extensionsResult } = await createAgentSession({',
	);
	text = text.replace(
		"\t\tmodel: playbookModel,\n\t\tsessionManager: SessionManager.inMemory(),",
		"\t\tmodel: playbookModel,\n\t\tmodelRuntime,\n\t\tsessionManager: SessionManager.inMemory(),",
	);
	text = text.replace(
		"        model: playbookModel,\n        sessionManager: SessionManager.inMemory(),",
		"        model: playbookModel,\n        modelRuntime,\n        sessionManager: SessionManager.inMemory(),",
	);
	text = text.replace(
		'\t// Override getApiKey to bypass real auth checks (on both agent and session)\n\t(session.agent as any).getApiKey = async () => "test-key";',
		'\tObject.defineProperty(session.agent, "streamFn", {\n\t\tconfigurable: true,\n\t\tget() { return (this as any).streamFunction; },\n\t\tset(value) { (this as any).streamFunction = value; },\n\t});\n\n\t// Override getApiKey to bypass real auth checks (on both agent and session)\n\t(session.agent as any).getApiKey = async () => "test-key";',
	);
	text = text.replace(
		'    // Override getApiKey to bypass real auth checks (on both agent and session)\n    session.agent.getApiKey = async () => "test-key";',
		'    Object.defineProperty(session.agent, "streamFn", {\n        configurable: true,\n        get() { return this.streamFunction; },\n        set(value) { this.streamFunction = value; },\n    });\n\n    // Override getApiKey to bypass real auth checks (on both agent and session)\n    session.agent.getApiKey = async () => "test-key";',
	);
	if (!text.includes("@earendil-works/pi-ai/compat")) {
		throw new Error(`${file}: expected compat getModel import`);
	}
	if (!text.includes("ModelRuntime")) {
		throw new Error(`${file}: expected ModelRuntime patch`);
	}
	if (!text.includes("streamFunction")) {
		throw new Error(`${file}: expected streamFn alias patch`);
	}
	writeFileSync(file, text);
}