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 starts a 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).
 */

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

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),
				},
			};
		}),
	};
}

// ---------------------------------------------------------------------------
// 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", () => {
		// Server is up
	});

	server.on("error", (err: NodeJS.ErrnoException) => {
		if (err.code === "EADDRINUSE") {
			console.warn(
				`[llama-proxy] Port ${PROXY_PORT} already in use — proxy not started. ` +
					`If a previous pi session left it running, kill it and reload.`,
			);
		} else {
			console.error("[llama-proxy] Server error:", err);
		}
	});

	return server;
}

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

export default function (pi: ExtensionAPI) {
	const 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`,
	});

	pi.on("session_end", async () => {
		server.close();
	});
}