added pi

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

@@ -0,0 +1,201 @@
+/**
+ * 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();
+	});
+}