rspace-online/server/cad-orchestrator.ts

/**
 * CAD Orchestrator — LLM-driven MCP tool calling for KiCad and FreeCAD
 *
 * Converts MCP tool schemas to Gemini function declarations, runs an agentic
 * loop where Gemini Flash plans and executes real MCP tool sequences, then
 * assembles artifacts (SVGs, exports) for the frontend.
 */

import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
import type { Tool } from "@modelcontextprotocol/sdk/types.js";

// ── MCP → Gemini schema conversion ──

/** Strip keys that Gemini function-calling schema rejects */
function cleanSchema(obj: any): any {
	if (obj === null || obj === undefined) return obj;
	if (Array.isArray(obj)) return obj.map(cleanSchema);
	if (typeof obj !== "object") return obj;

	const out: any = {};
	for (const [k, v] of Object.entries(obj)) {
		// Gemini rejects these in FunctionDeclaration parameters
		if (["default", "additionalProperties", "$schema", "examples", "title"].includes(k)) continue;
		out[k] = cleanSchema(v);
	}
	return out;
}

/** Convert a single MCP Tool to a Gemini FunctionDeclaration */
export function mcpToolToGeminiFn(tool: Tool): any {
	const params = tool.inputSchema ? cleanSchema(tool.inputSchema) : { type: "object", properties: {} };
	// Gemini requires type:"OBJECT" (uppercase) in some SDK versions, but @google/generative-ai handles lowercase
	return {
		name: tool.name,
		description: tool.description || tool.name,
		parameters: params,
	};
}

// ── Extract text from MCP CallToolResult ──

export function extractMcpText(result: any): string {
	const content = result?.content;
	if (!Array.isArray(content)) return JSON.stringify(result);
	return content
		.filter((c: any) => c.type === "text")
		.map((c: any) => c.text || "")
		.join("\n");
}

// ── Agentic loop ──

export interface ToolCallLogEntry {
	tool: string;
	args: Record<string, any>;
	result: string;
}

export interface OrchestrationResult {
	artifacts: Map<string, string>;
	finalMessage: string;
	toolCallLog: ToolCallLogEntry[];
}

export async function runCadAgentLoop(
	client: Client,
	systemPrompt: string,
	userPrompt: string,
	geminiApiKey: string,
	maxTurns = 8,
): Promise<OrchestrationResult> {
	// 1. Fetch real tool schemas from MCP server
	const { tools } = await client.listTools();
	const functionDeclarations = tools.map(mcpToolToGeminiFn);

	// 2. Initialize Gemini Flash with function calling
	const { GoogleGenerativeAI } = await import("@google/generative-ai");
	const genAI = new GoogleGenerativeAI(geminiApiKey);
	const model = genAI.getGenerativeModel({
		model: "gemini-2.5-flash",
		systemInstruction: systemPrompt,
		generationConfig: { temperature: 0.2 },
		tools: [{ functionDeclarations }],
	});

	// 3. Agentic loop
	const artifacts = new Map<string, string>();
	const toolCallLog: ToolCallLogEntry[] = [];
	const deadline = Date.now() + 60_000;

	let contents: any[] = [
		{ role: "user", parts: [{ text: userPrompt }] },
	];

	for (let turn = 0; turn < maxTurns; turn++) {
		if (Date.now() > deadline) {
			console.warn("[cad-orchestrator] 60s deadline reached");
			break;
		}

		const result = await model.generateContent({ contents });
		const candidate = result.response.candidates?.[0];
		if (!candidate) break;

		const parts = candidate.content?.parts || [];
		const fnCalls = parts.filter((p: any) => p.functionCall);

		if (fnCalls.length === 0) {
			// Done — extract final text
			const finalText = parts
				.filter((p: any) => p.text)
				.map((p: any) => p.text)
				.join("\n");
			return { artifacts, finalMessage: finalText || "Design complete.", toolCallLog };
		}

		// Execute each function call on MCP
		const fnResponseParts: any[] = [];
		for (const part of fnCalls) {
			const fc = part.functionCall!;
			console.log(`[cad-orchestrator] Turn ${turn + 1}: ${fc.name}(${JSON.stringify(fc.args).slice(0, 200)})`);

			let mcpResultText: string;
			try {
				const mcpResult = await client.callTool({
					name: fc.name,
					arguments: (fc.args || {}) as Record<string, unknown>,
				});
				mcpResultText = extractMcpText(mcpResult);
			} catch (err) {
				mcpResultText = `Error: ${err instanceof Error ? err.message : String(err)}`;
			}

			toolCallLog.push({ tool: fc.name, args: fc.args || {}, result: mcpResultText });
			artifacts.set(fc.name, mcpResultText);

			fnResponseParts.push({
				functionResponse: {
					name: fc.name,
					response: { result: mcpResultText },
				},
			});
		}

		// Feed results back for next turn
		contents.push({ role: "model", parts });
		contents.push({ role: "user", parts: fnResponseParts });
	}

	return {
		artifacts,
		finalMessage: "Design generation completed (max turns reached).",
		toolCallLog,
	};
}

// ── Result assemblers ──

export interface KicadResult {
	schematicSvg: string | null;
	boardSvg: string | null;
	gerberUrl: string | null;
	bomUrl: string | null;
	pdfUrl: string | null;
	drcResults: { violations: string[] } | null;
	summary: string;
	toolCallLog: ToolCallLogEntry[];
}

export function assembleKicadResult(orch: OrchestrationResult): KicadResult {
	let schematicSvg: string | null = null;
	let boardSvg: string | null = null;
	let gerberUrl: string | null = null;
	let bomUrl: string | null = null;
	let pdfUrl: string | null = null;
	let drcResults: { violations: string[] } | null = null;

	for (const entry of orch.toolCallLog) {
		try {
			const parsed = JSON.parse(entry.result);
			switch (entry.tool) {
				case "export_svg":
					// Could be schematic or board SVG — check args or content
					if (entry.args.type === "board" || entry.args.board) {
						boardSvg = parsed.svg_path || parsed.path || parsed.url || null;
					} else {
						schematicSvg = parsed.svg_path || parsed.path || parsed.url || null;
					}
					break;
				case "run_drc":
					drcResults = {
						violations: parsed.violations || parsed.errors || [],
					};
					break;
				case "export_gerber":
					gerberUrl = parsed.gerber_path || parsed.path || parsed.url || null;
					break;
				case "export_bom":
					bomUrl = parsed.bom_path || parsed.path || parsed.url || null;
					break;
				case "export_pdf":
					pdfUrl = parsed.pdf_path || parsed.path || parsed.url || null;
					break;
			}
		} catch {
			// Non-JSON results are fine (intermediate steps)
		}
	}

	return {
		schematicSvg,
		boardSvg,
		gerberUrl,
		bomUrl,
		pdfUrl,
		drcResults,
		summary: orch.finalMessage,
		toolCallLog: orch.toolCallLog,
	};
}

export interface FreecadResult {
	previewUrl: string | null;
	stepUrl: string | null;
	stlUrl: string | null;
	summary: string;
	toolCallLog: ToolCallLogEntry[];
}

export function assembleFreecadResult(orch: OrchestrationResult): FreecadResult {
	let previewUrl: string | null = null;
	let stepUrl: string | null = null;
	let stlUrl: string | null = null;

	for (const entry of orch.toolCallLog) {
		try {
			const parsed = JSON.parse(entry.result);
			// FreeCAD exports via execute_python_script — look for file paths in results
			if (entry.tool === "execute_python_script" || entry.tool === "execute_script") {
				const text = entry.result.toLowerCase();
				if (text.includes(".step") || text.includes(".stp")) {
					stepUrl = parsed.path || parsed.file_path || extractPathFromText(entry.result, [".step", ".stp"]);
				}
				if (text.includes(".stl")) {
					stlUrl = parsed.path || parsed.file_path || extractPathFromText(entry.result, [".stl"]);
				}
			}
			// save_document may also produce a path
			if (entry.tool === "save_document") {
				const path = parsed.path || parsed.file_path || null;
				if (path && (path.endsWith(".FCStd") || path.endsWith(".fcstd"))) {
					// Not directly servable, but note it
				}
			}
		} catch {
			// Try extracting paths from raw text
			stepUrl = stepUrl || extractPathFromText(entry.result, [".step", ".stp"]);
			stlUrl = stlUrl || extractPathFromText(entry.result, [".stl"]);
		}
	}

	return {
		previewUrl,
		stepUrl,
		stlUrl,
		summary: orch.finalMessage,
		toolCallLog: orch.toolCallLog,
	};
}

/** Extract a file path ending with one of the given extensions from text */
function extractPathFromText(text: string, extensions: string[]): string | null {
	for (const ext of extensions) {
		const regex = new RegExp(`(/[\\w./-]+${ext.replace(".", "\\.")})`, "i");
		const match = text.match(regex);
		if (match) return match[1];
	}
	return null;
}

// ── System prompts ──

export const KICAD_SYSTEM_PROMPT = `You are a KiCad PCB design assistant. You have access to KiCad MCP tools to create real PCB designs.

Follow this workflow:
1. create_project — Create a new KiCad project in /data/files/generated/kicad-<timestamp>/
2. search_symbols — Find component symbols in KiCad libraries (e.g. ESP32, BME280, capacitors, resistors)
3. add_schematic_component — Place each component on the schematic
4. add_schematic_net_label — Add net labels for connections
5. add_schematic_connection — Wire components together
6. generate_netlist — Generate the netlist from schematic
7. place_component — Place footprints on the board
8. route_trace — Route traces between pads
9. export_svg — Export schematic SVG (type: "schematic") and board SVG (type: "board")
10. run_drc — Run Design Rule Check
11. export_gerber, export_bom, export_pdf — Generate manufacturing outputs

Important:
- Use /data/files/generated/kicad-${Date.now()}/ as the project directory
- Search for real symbols before placing components
- Add decoupling capacitors and pull-up resistors as needed
- Set reasonable board outline dimensions
- After placing components, route critical traces
- Always run DRC before exporting
- If a tool call fails, try an alternative approach rather than repeating the same call`;

export const FREECAD_SYSTEM_PROMPT = `You are a FreeCAD parametric CAD assistant. You have access to FreeCAD MCP tools to create real 3D models.

Follow this workflow:
1. execute_python_script — Create output directory: import os; os.makedirs("/data/files/generated/freecad-<timestamp>", exist_ok=True)
2. Create base geometry using create_box, create_cylinder, or create_sphere
3. Use boolean_operation (union, cut, intersection) to combine shapes
4. list_objects to verify the model state
5. save_document to save the FreeCAD file
6. execute_python_script to export STEP: Part.export([obj], "/data/files/generated/freecad-<id>/model.step")
7. execute_python_script to export STL: Mesh.export([obj], "/data/files/generated/freecad-<id>/model.stl")

Important:
- Use /data/files/generated/freecad-${Date.now()}/ as the working directory
- For hollow objects, create the outer shell then cut the inner volume
- For complex shapes, build up from primitives with boolean operations
- Wall thickness should be at least 1mm for 3D printing
- Always export both STEP (for CAD) and STL (for 3D printing)
- If a tool call fails, try an alternative approach`;