dotfiles/pi/.pi/agent/extensions/llama-schema-proxy.ts

/**
 * llama-server Schema Sanitization Proxy
 *
 * llama-server strictly validates JSON Schema and rejects any schema node
 * that lacks a `type` field. Some of pi's built-in tools (e.g. `subagent`)
 * have complex union-type parameters represented as `{"description": "..."}` with
 * no `type`, which causes llama-server to return a 400 error.
 *
 * This extension provides an optional tiny local HTTP proxy on port 8081 that:
 *   1. Intercepts outgoing OpenAI-compatible API calls
 *   2. Walks tool schemas and adds `"type": "string"` to any schema node
 *      that is missing a type declaration
 *   3. Forwards the fixed request to llama-server on port 8080
 *   4. Streams the response back transparently
 *
 * It also overrides the `llama-cpp` provider's baseUrl to point at the proxy,
 * so no changes to models.json are needed (beyond what's already there).
 *
 * Use `/llama-proxy` command to toggle the proxy on/off. Off by default.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import * as http from "http";
import { execSync } from "child_process";

const PROXY_PORT = 8081;
const TARGET_HOST = "127.0.0.1";
const TARGET_PORT = 8080;

// ---------------------------------------------------------------------------
// Schema sanitizer
// ---------------------------------------------------------------------------

/**
 * Recursively walk a JSON Schema object and add `"type": "string"` to any
 * node that has no `type` and no composition keywords (oneOf/anyOf/allOf/$ref).
 * This satisfies llama-server's strict validation without breaking valid nodes.
 */
function sanitizeSchema(schema: unknown): unknown {
	if (!schema || typeof schema !== "object") return schema;
	if (Array.isArray(schema)) return schema.map(sanitizeSchema);

	const obj = schema as Record<string, unknown>;
	const result: Record<string, unknown> = {};

	for (const [key, value] of Object.entries(obj)) {
		if (key === "properties" && value && typeof value === "object" && !Array.isArray(value)) {
			result[key] = Object.fromEntries(
				Object.entries(value as Record<string, unknown>).map(([k, v]) => [k, sanitizeSchema(v)]),
			);
		} else if (key === "items") {
			result[key] = sanitizeSchema(value);
		} else if (key === "additionalProperties" && value && typeof value === "object") {
			result[key] = sanitizeSchema(value);
		} else if (
			(key === "oneOf" || key === "anyOf" || key === "allOf") &&
			Array.isArray(value)
		) {
			result[key] = value.map(sanitizeSchema);
		} else {
			result[key] = value;
		}
	}

	// If this schema node has no type and no composition keywords, default to "string"
	const hasType = "type" in result;
	const hasComposition =
		"oneOf" in result || "anyOf" in result || "allOf" in result || "$ref" in result;
	const hasEnum = "enum" in result || "const" in result;

	if (!hasType && !hasComposition && !hasEnum) {
		result["type"] = "string";
	}

	return result;
}

/**
 * Patch the `tools` array in a parsed request body, if present.
 */
function sanitizeRequestBody(body: Record<string, unknown>): Record<string, unknown> {
	if (!Array.isArray(body.tools)) return body;

	return {
		...body,
		tools: (body.tools as unknown[]).map((tool) => {
			if (!tool || typeof tool !== "object") return tool;
			const t = tool as Record<string, unknown>;
			if (!t.function || typeof t.function !== "object") return t;
			const fn = t.function as Record<string, unknown>;
			if (!fn.parameters) return t;
			return {
				...t,
				function: {
					...fn,
					parameters: sanitizeSchema(fn.parameters),
				},
			};
		}),
	};
}

// ---------------------------------------------------------------------------
// Process management
// ---------------------------------------------------------------------------

/**
 * Kill any existing processes using the proxy port.
 */
function killExistingProxy(): void {
	try {
		// Use lsof to find processes on the port and kill them
		const output = execSync(`lsof -ti:${PROXY_PORT} 2>/dev/null || true`, {
			encoding: "utf-8",
		});
		const pids = output.trim().split("\n").filter(Boolean);
		for (const pid of pids) {
			try {
				process.kill(Number(pid), "SIGTERM");
				console.log(`[llama-proxy] Terminated old instance (PID: ${pid})`);
			} catch {
				// Process may have already exited
			}
		}
	} catch {
		// lsof not available or other error — continue anyway
	}
}

// ---------------------------------------------------------------------------
// Proxy server
// ---------------------------------------------------------------------------

function startProxy(): http.Server {
	const server = http.createServer((req, res) => {
		const chunks: Buffer[] = [];

		req.on("data", (chunk: Buffer) => chunks.push(chunk));
		req.on("end", () => {
			const rawBody = Buffer.concat(chunks).toString("utf-8");

			// Attempt to sanitize schemas in JSON bodies
			let forwardBody = rawBody;
			const contentType = req.headers["content-type"] ?? "";
			if (contentType.includes("application/json") && rawBody.trim().startsWith("{")) {
				try {
					const parsed = JSON.parse(rawBody) as Record<string, unknown>;
					const sanitized = sanitizeRequestBody(parsed);
					forwardBody = JSON.stringify(sanitized);
				} catch {
					// Not valid JSON — send as-is
				}
			}

			const forwardBuffer = Buffer.from(forwardBody, "utf-8");

			// Build forwarded headers, updating host and content-length
			const forwardHeaders: Record<string, string | string[]> = {};
			for (const [k, v] of Object.entries(req.headers)) {
				if (k === "host") continue; // rewrite below
				if (v !== undefined) forwardHeaders[k] = v as string | string[];
			}
			forwardHeaders["host"] = `${TARGET_HOST}:${TARGET_PORT}`;
			forwardHeaders["content-length"] = String(forwardBuffer.byteLength);

			const proxyReq = http.request(
				{
					host: TARGET_HOST,
					port: TARGET_PORT,
					path: req.url,
					method: req.method,
					headers: forwardHeaders,
				},
				(proxyRes) => {
					res.writeHead(proxyRes.statusCode ?? 200, proxyRes.headers);
					proxyRes.pipe(res, { end: true });
				},
			);

			proxyReq.on("error", (err) => {
				const msg = `Proxy error forwarding to llama-server: ${err.message}`;
				if (!res.headersSent) {
					res.writeHead(502, { "content-type": "text/plain" });
				}
				res.end(msg);
			});

			proxyReq.write(forwardBuffer);
			proxyReq.end();
		});

		req.on("error", (err) => {
			console.error("[llama-proxy] request error:", err);
		});
	});

	server.listen(PROXY_PORT, "127.0.0.1", () => {
		console.log(`[llama-proxy] Proxy started on port ${PROXY_PORT}`);
	});

	server.on("error", (err: NodeJS.ErrnoException) => {
		if (err.code === "EADDRINUSE") {
			console.error(
				`[llama-proxy] Port ${PROXY_PORT} already in use. ` +
					`Killing old instances and retrying...`,
			);
			killExistingProxy();
		} else {
			console.error("[llama-proxy] Server error:", err);
		}
	});

	return server;
}

// ---------------------------------------------------------------------------
// Extension entry point
// ---------------------------------------------------------------------------

export default function (pi: ExtensionAPI) {
	let server: http.Server | null = null;
	let proxyEnabled = false;

	/**
	 * Start the proxy and register the provider override.
	 */
	function enableProxy(): void {
		if (proxyEnabled) {
			console.log("[llama-proxy] Proxy already enabled");
			return;
		}

		killExistingProxy();
		server = startProxy();

		// Override the llama-cpp provider's baseUrl to route through our proxy.
		// models.json model definitions are preserved; only the endpoint changes.
		pi.registerProvider("llama-cpp", {
			baseUrl: `http://127.0.0.1:${PROXY_PORT}/v1`,
		});

		proxyEnabled = true;
		console.log("[llama-proxy] Proxy enabled");
	}

	/**
	 * Disable the proxy and restore default provider.
	 */
	function disableProxy(): void {
		if (!proxyEnabled) {
			console.log("[llama-proxy] Proxy already disabled");
			return;
		}

		if (server) {
			server.close();
			server = null;
		}

		// Reset provider to default (no baseUrl override)
		pi.registerProvider("llama-cpp", {});

		proxyEnabled = false;
		console.log("[llama-proxy] Proxy disabled");
	}

	// Register the /llama-proxy command to toggle the proxy
	pi.registerCommand("llama-proxy", async (args) => {
		const action = args[0]?.toLowerCase() || "";

		if (action === "on") {
			enableProxy();
		} else if (action === "off") {
			disableProxy();
		} else if (action === "status") {
			console.log(`[llama-proxy] Status: ${proxyEnabled ? "enabled" : "disabled"}`);
		} else {
			// Toggle if no argument
			if (proxyEnabled) {
				disableProxy();
			} else {
				enableProxy();
			}
		}
	});

	// Clean up on session end
	pi.on("session_end", async () => {
		if (server) {
			server.close();
		}
	});
}