Mechanical Engineering Agent — PRD & Architecture

// lib/llm/types.ts — Unified types across providers
interface ToolDefinition {
  name: string;
  description: string;
  parameters: JSONSchema;
}

interface ToolCall {
  id: string;
  name: string;
  arguments: Record<string, unknown>;
}

interface ToolResult {
  tool_call_id: string;
  content: string;
  is_error?: boolean;
}

interface Message {
  role: 'system' | 'user' | 'assistant' | 'tool';
  content: string | ContentBlock[];
  tool_calls?: ToolCall[];
  tool_results?: ToolResult[];
}

interface CompletionOptions {
  model?: string;
  tools?: ToolDefinition[];
  temperature?: number;
  maxTokens?: number;
  responseFormat?: { type: 'json_object' | 'text' };
  stop?: string[];
}

interface LLMResponse {
  content: string;
  tool_calls?: ToolCall[];
  usage: { prompt_tokens: number; completion_tokens: number };
  model: string;
  finish_reason: 'stop' | 'tool_use' | 'length';
}

interface LLMChunk {
  content?: string;
  tool_call?: Partial<ToolCall>;
  finish_reason?: string;
}

// lib/llm/provider.ts — Provider interface
interface LLMProvider {
  name: string;
  complete(messages: Message[], options?: CompletionOptions): Promise<LLMResponse>;
  stream(messages: Message[], options?: CompletionOptions): AsyncIterable<LLMChunk>;
}

// lib/llm/anthropic-provider.ts
import Anthropic from '@anthropic-ai/sdk';

class AnthropicProvider implements LLMProvider {
  name = 'anthropic';
  private client: Anthropic;

  constructor(config: { apiKey: string }) {
    this.client = new Anthropic({ apiKey: config.apiKey });
  }

  async complete(messages: Message[], options?: CompletionOptions): Promise<LLMResponse> {
    const response = await this.client.messages.create({
      model: options?.model ?? 'claude-sonnet-4-5-20250929',
      max_tokens: options?.maxTokens ?? 4096,
      temperature: options?.temperature ?? 0.3,
      system: this.extractSystem(messages),
      messages: this.mapMessages(messages),
      tools: options?.tools?.map(t => this.mapTool(t)),
    });

    return this.mapResponse(response);
  }

  // Maps MetaForge ToolDefinition → Anthropic tool_use format
  private mapTool(tool: ToolDefinition): Anthropic.Tool {
    return {
      name: tool.name,
      description: tool.description,
      input_schema: tool.parameters,
    };
  }

  async *stream(messages: Message[], options?: CompletionOptions): AsyncIterable<LLMChunk> {
    const stream = this.client.messages.stream({
      model: options?.model ?? 'claude-sonnet-4-5-20250929',
      max_tokens: options?.maxTokens ?? 4096,
      messages: this.mapMessages(messages),
    });

    for await (const event of stream) {
      yield this.mapChunk(event);
    }
  }
}

// lib/llm/openai-provider.ts
import OpenAI from 'openai';

class OpenAIProvider implements LLMProvider {
  name = 'openai';
  private client: OpenAI;

  constructor(config: { apiKey: string }) {
    this.client = new OpenAI({ apiKey: config.apiKey });
  }

  async complete(messages: Message[], options?: CompletionOptions): Promise<LLMResponse> {
    const response = await this.client.chat.completions.create({
      model: options?.model ?? 'gpt-4o',
      temperature: options?.temperature ?? 0.3,
      max_tokens: options?.maxTokens ?? 4096,
      messages: this.mapMessages(messages),
      tools: options?.tools?.map(t => this.mapTool(t)),
      response_format: options?.responseFormat,
    });

    return this.mapResponse(response);
  }

  // Maps MetaForge ToolDefinition → OpenAI function_calling format
  private mapTool(tool: ToolDefinition): OpenAI.ChatCompletionTool {
    return {
      type: 'function',
      function: {
        name: tool.name,
        description: tool.description,
        parameters: tool.parameters,
      },
    };
  }
}

// lib/llm/factory.ts
type ProviderType = 'anthropic' | 'openai';

interface LLMConfig {
  provider: ProviderType;
  apiKey: string;
  defaultModel?: string;
}

class LLMProviderFactory {
  static create(config: LLMConfig): LLMProvider {
    switch (config.provider) {
      case 'anthropic':
        return new AnthropicProvider({ apiKey: config.apiKey });
      case 'openai':
        return new OpenAIProvider({ apiKey: config.apiKey });
      default:
        throw new Error(`Unknown LLM provider: ${config.provider}`);
    }
  }
}

// agents/mechanical/mechanical-agent.ts
import { Agent, AgentConfig, ExecutionContext, AgentResult } from '@metaforge/core';
import { MechanicalInputSchema, MechanicalOutputSchema } from './schemas';

class MechanicalAgent implements Agent {
  name = 'mechanical-agent';
  version = '1.0.0';
  description = 'Mechanical engineering: CAD generation, FEA, CFD, tolerance analysis, material selection';

  private llm: LLMProvider;
  private tools: ToolRegistry;

  async initialize(config: AgentConfig): Promise<void> {
    this.llm = config.llmProvider;
    this.tools = config.toolRegistry;
  }

  registerTools(registry: ToolRegistry): void {
    registry.register(new FreeCADAdapter());
    registry.register(new CalculiXAdapter());
    registry.register(new OpenFOAMAdapter());
  }

  async execute(context: ExecutionContext): Promise<AgentResult> {
    const input = MechanicalInputSchema.parse(context.input);
    const taskType = this.classifyTask(input);

    switch (taskType) {
      case 'enclosure_generation':
        return await this.generateEnclosure(input);
      case 'fea_stress':
        return await this.runStressFEA(input);
      case 'fea_modal':
        return await this.runModalFEA(input);
      case 'fea_thermal':
        return await this.runThermalFEA(input);
      case 'cfd_thermal':
        return await this.runCFDThermal(input);
      case 'tolerance_analysis':
        return await this.runToleranceAnalysis(input);
      case 'material_selection':
        return await this.selectMaterial(input);
      case 'mechanical_bom':
        return await this.generateMechanicalBOM(input);
      case 'dfa_analysis':
        return await this.runDFAAnalysis(input);
      default:
        throw new Error(`Unknown task type: ${taskType}`);
    }
  }

  validateInput(input: unknown): ValidationResult {
    const result = MechanicalInputSchema.safeParse(input);
    return result.success
      ? { valid: true }
      : { valid: false, errors: result.error.issues };
  }

  validateOutput(output: unknown): ValidationResult {
    const result = MechanicalOutputSchema.safeParse(output);
    return result.success
      ? { valid: true }
      : { valid: false, errors: result.error.issues };
  }
}

interface FreeCADAdapter extends ToolAdapter {
  name: 'freecad';
  version: string;
  capabilities: ['parametric_modeling', 'step_export', 'stl_export', 'mesh_generation', 'boolean_ops'];

  // Detection
  detect(): Promise<ToolInstall | null>;

  // Script execution (headless)
  executeScript(script: string, args?: Record<string, string>): Promise<FreeCADResult>;

  // Export operations
  exportSTEP(modelPath: string, outputPath: string): Promise<string>;
  exportSTL(modelPath: string, outputPath: string, meshParams?: MeshParams): Promise<string>;

  // Mesh generation (for FEA)
  generateMesh(modelPath: string, params: MeshParams): Promise<MeshResult>;

  // Model inspection
  getModelInfo(modelPath: string): Promise<ModelInfo>;
  getBoundingBox(modelPath: string): Promise<BoundingBox>;

  // Health
  healthCheck(): Promise<HealthStatus>;
}

interface FreeCADResult {
  exitCode: number;
  stdout: string;
  stderr: string;
  outputFiles: string[];       // paths to generated files
  modelInfo?: ModelInfo;
}

interface MeshParams {
  algorithm: 'gmsh' | 'netgen';
  maxElementSize: number;      // mm
  minElementSize: number;      // mm
  order: 1 | 2;               // linear or quadratic elements
}

interface MeshResult {
  meshFile: string;            // path to .inp or .unv mesh file
  nodeCount: number;
  elementCount: number;
  qualityMetrics: {
    minAngle: number;
    maxAspectRatio: number;
    jacobianRatio: number;
  };
}

interface ModelInfo {
  volume: number;              // mm³
  surfaceArea: number;         // mm²
  boundingBox: BoundingBox;
  centerOfMass: [number, number, number];
  bodyCount: number;
}

interface BoundingBox {
  min: [number, number, number];
  max: [number, number, number];
  dimensions: [number, number, number];  // width, depth, height in mm
}

class FreeCADAdapter implements ToolAdapter {
  name = 'freecad';
  version = '1.0.0';

  async detect(): Promise<ToolInstall | null> {
    const searchPaths = [
      '/usr/bin/freecadcmd',
      '/usr/local/bin/freecadcmd',
      '/snap/bin/freecadcmd',
      '/Applications/FreeCAD.app/Contents/MacOS/FreeCADCmd',
      'C:\\Program Files\\FreeCAD 0.21\\bin\\FreeCADCmd.exe',
    ];

    for (const path of searchPaths) {
      try {
        const result = await exec(`${path} --version`);
        const version = this.parseVersion(result.stdout);
        if (semver.gte(version, '0.21.0')) {
          return { found: true, version, path };
        }
      } catch { /* not found at this path */ }
    }
    return null;
  }

  async executeScript(script: string, args?: Record<string, string>): Promise<FreeCADResult> {
    const scriptPath = await this.writeScript(script);      // temp file
    const argStr = args
      ? Object.entries(args).map(([k, v]) => `--${k} ${v}`).join(' ')
      : '';

    const result = await exec(
      `freecadcmd ${scriptPath} -- ${argStr}`,
      { timeout: 120_000, cwd: this.workDir }               // 2 min timeout
    );

    return {
      exitCode: result.exitCode,
      stdout: result.stdout,
      stderr: result.stderr,
      outputFiles: this.findOutputFiles(this.workDir),
    };
  }

  async exportSTEP(modelPath: string, outputPath: string): Promise<string> {
    const script = `
import FreeCAD
import Part
doc = FreeCAD.openDocument("${modelPath}")
Part.export(doc.Objects, "${outputPath}")
doc.close()
    `;
    await this.executeScript(script);
    return outputPath;
  }

  async generateMesh(modelPath: string, params: MeshParams): Promise<MeshResult> {
    const script = `
import FreeCAD
import Fem
import femmesh.femmesh2mesh

doc = FreeCAD.openDocument("${modelPath}")
shape = doc.Objects[0].Shape

femmesh = doc.addObject("Fem::FemMeshShapeNetgenObject", "FEMMesh")
femmesh.Shape = doc.Objects[0]
femmesh.MaxSize = ${params.maxElementSize}
femmesh.MinSize = ${params.minElementSize}
femmesh.SecondOrder = ${params.order === 2 ? 'True' : 'False'}
doc.recompute()

# Export as .inp for CalculiX
import feminout.importInpMesh
feminout.importInpMesh.export([femmesh], "${modelPath.replace('.FCStd', '.inp')}")
doc.close()
    `;

    const result = await this.executeScript(script);
    const meshFile = modelPath.replace('.FCStd', '.inp');

    return {
      meshFile,
      nodeCount: this.extractNodeCount(result.stdout),
      elementCount: this.extractElementCount(result.stdout),
      qualityMetrics: this.extractMeshQuality(result.stdout),
    };
  }
}

interface CalculiXAdapter extends ToolAdapter {
  name: 'calculix';
  version: string;
  capabilities: ['static_stress', 'modal_analysis', 'thermal_steady', 'thermal_transient'];

  // Detection
  detect(): Promise<ToolInstall | null>;

  // FEA operations
  runStaticStress(config: StaticStressConfig): Promise<StressResult>;
  runModalAnalysis(config: ModalConfig): Promise<ModalResult>;
  runThermalAnalysis(config: ThermalConfig): Promise<ThermalResult>;

  // Low-level
  execute(inpFile: string): Promise<CalculiXResult>;
  parseResults(frdFile: string): Promise<FEAFieldData>;
}

interface StaticStressConfig {
  meshFile: string;            // .inp mesh file from FreeCAD
  material: MaterialProperties;
  loads: Load[];
  constraints: BoundaryCondition[];
}

interface StressResult {
  maxVonMises: number;         // MPa
  maxDisplacement: number;     // mm
  safetyFactor: number;        // yield_strength / max_stress
  stressField: string;         // path to .frd results file
  criticalLocations: CriticalLocation[];
  passed: boolean;             // safetyFactor >= minSafetyFactor
}

interface ModalConfig {
  meshFile: string;
  material: MaterialProperties;
  constraints: BoundaryCondition[];
  numModes: number;            // number of modes to extract (typically 6-10)
}

interface ModalResult {
  frequencies: number[];       // Hz, natural frequencies
  modeShapes: string;          // path to .frd mode shape data
  participationFactors: number[];
  passed: boolean;             // no modes within excitation frequency range
}

interface ThermalConfig {
  meshFile: string;
  material: MaterialProperties;
  heatSources: HeatSource[];
  convection: ConvectionBC[];
  radiation?: RadiationBC[];
  analysisType: 'steady_state' | 'transient';
  timeStep?: number;           // seconds (for transient)
  totalTime?: number;          // seconds (for transient)
}

interface ThermalResult {
  maxTemperature: number;      // °C
  minTemperature: number;      // °C
  temperatureField: string;    // path to .frd results
  heatFlux: number;            // W/m²
  thermalResistance: number;   // °C/W
  hotspots: Hotspot[];
  passed: boolean;             // maxTemp <= maxAllowable
}

interface Load {
  type: 'force' | 'pressure' | 'gravity' | 'centrifugal';
  magnitude: number;
  direction?: [number, number, number];
  surface?: string;            // surface set name
}

interface BoundaryCondition {
  type: 'fixed' | 'displacement' | 'symmetry';
  surface: string;             // surface set name
  values?: [number, number, number];
}

interface HeatSource {
  type: 'volume' | 'surface';
  power: number;               // Watts
  location: string;            // body or surface set name
}

interface ConvectionBC {
  surface: string;
  htc: number;                 // W/(m²·K), heat transfer coefficient
  ambientTemp: number;         // °C
}

interface Hotspot {
  location: [number, number, number];
  temperature: number;         // °C
  nearestComponent?: string;
}

class CalculiXAdapter implements ToolAdapter {
  name = 'calculix';
  version = '1.0.0';

  async detect(): Promise<ToolInstall | null> {
    try {
      const result = await exec('ccx -v');
      const version = this.parseVersion(result.stdout);
      return { found: true, version, path: await which('ccx') };
    } catch {
      return null;
    }
  }

  async runStaticStress(config: StaticStressConfig): Promise<StressResult> {
    // 1. Generate .inp file with material, loads, BCs
    const inpContent = this.buildStaticInp(config);
    const inpPath = path.join(this.workDir, 'stress_analysis.inp');
    await writeFile(inpPath, inpContent);

    // 2. Execute CalculiX
    const result = await exec(`ccx -i ${inpPath.replace('.inp', '')}`, {
      timeout: 300_000,  // 5 min timeout
      cwd: this.workDir,
    });

    // 3. Parse .frd results file
    const frdPath = inpPath.replace('.inp', '.frd');
    const fieldData = await this.parseResults(frdPath);

    return {
      maxVonMises: fieldData.maxVonMises,
      maxDisplacement: fieldData.maxDisplacement,
      safetyFactor: config.material.yieldStrength / fieldData.maxVonMises,
      stressField: frdPath,
      criticalLocations: fieldData.criticalLocations,
      passed: (config.material.yieldStrength / fieldData.maxVonMises) >= 1.5,
    };
  }

  private buildStaticInp(config: StaticStressConfig): string {
    return `
*HEADING
MetaForge ME Agent - Static Stress Analysis
*INCLUDE, INPUT=${config.meshFile}
*MATERIAL, NAME=MATERIAL1
*ELASTIC
${config.material.youngsModulus}, ${config.material.poissonsRatio}
*DENSITY
${config.material.density}
*SOLID SECTION, ELSET=EALL, MATERIAL=MATERIAL1
*STEP
*STATIC
*BOUNDARY
${this.formatBoundaryConditions(config.constraints)}
*CLOAD
${this.formatLoads(config.loads)}
*NODE FILE
U
*EL FILE
S
*END STEP
    `.trim();
  }
}

interface OpenFOAMAdapter extends ToolAdapter {
  name: 'openfoam';
  version: string;
  capabilities: ['cfd_thermal', 'flow_simulation', 'mesh_generation'];

  // Detection
  detect(): Promise<ToolInstall | null>;

  // Case management
  setupCase(config: CFDCaseConfig): Promise<string>;  // returns case directory path

  // Meshing
  generateMesh(caseDir: string, geometry: string): Promise<OpenFOAMMeshResult>;

  // Solving
  runSolver(caseDir: string, solver: OpenFOAMSolver): Promise<CFDResult>;

  // Post-processing
  extractResults(caseDir: string): Promise<CFDFieldData>;
}

type OpenFOAMSolver = 'buoyantSimpleFoam' | 'simpleFoam' | 'buoyantPimpleFoam';

interface CFDCaseConfig {
  solver: OpenFOAMSolver;
  geometry: {
    stlFile: string;           // enclosure geometry as STL
    inlets?: Surface[];
    outlets?: Surface[];
    walls?: Surface[];
  };
  thermalSources: ThermalSource[];
  fluidProperties: FluidProperties;
  meshSettings: {
    baseCellSize: number;      // mm
    refinementLevels: number;
    boundaryLayers: number;
  };
}

interface ThermalSource {
  name: string;
  type: 'fixedTemperature' | 'fixedHeatFlux';
  value: number;               // °C or W/m²
  surface: string;             // patch name
}

interface FluidProperties {
  fluid: 'air' | 'water' | 'custom';
  density?: number;            // kg/m³
  viscosity?: number;          // Pa·s
  thermalConductivity?: number; // W/(m·K)
  specificHeat?: number;       // J/(kg·K)
  beta?: number;               // 1/K, thermal expansion coefficient
}

interface CFDResult {
  converged: boolean;
  iterations: number;
  residuals: { [field: string]: number };
  maxVelocity: number;         // m/s
  maxTemperature: number;      // °C
  avgTemperature: number;      // °C
  heatTransferRate: number;    // W
  pressureDrop?: number;       // Pa
  results: CFDFieldData;
}

interface CFDFieldData {
  temperatureField: string;    // path to OpenFOAM field data
  velocityField: string;
  pressureField: string;
  hotspots: Hotspot[];
}

class OpenFOAMAdapter implements ToolAdapter {
  name = 'openfoam';
  version = '1.0.0';

  async detect(): Promise<ToolInstall | null> {
    try {
      const result = await exec('simpleFoam -help');
      // OpenFOAM prints version in the banner
      const version = this.parseVersion(result.stdout);
      return { found: true, version, path: process.env.WM_PROJECT_DIR || '' };
    } catch {
      return null;
    }
  }

  async setupCase(config: CFDCaseConfig): Promise<string> {
    const caseDir = path.join(this.workDir, 'cfd_case');
    await mkdir(caseDir, { recursive: true });

    // Create OpenFOAM directory structure
    await mkdir(path.join(caseDir, '0'));        // initial conditions
    await mkdir(path.join(caseDir, 'constant')); // mesh + properties
    await mkdir(path.join(caseDir, 'system'));    // solver settings

    // Write dictionary files
    await this.writeTransportProperties(caseDir, config.fluidProperties);
    await this.writeFvSchemes(caseDir);
    await this.writeFvSolution(caseDir, config.solver);
    await this.writeControlDict(caseDir, config.solver);
    await this.writeInitialConditions(caseDir, config);
    await this.writeSnappyHexMeshDict(caseDir, config);

    return caseDir;
  }

  async generateMesh(caseDir: string, stlFile: string): Promise<OpenFOAMMeshResult> {
    // 1. Block mesh (background mesh)
    await exec('blockMesh', { cwd: caseDir, timeout: 60_000 });

    // 2. Copy STL to triSurface
    const triDir = path.join(caseDir, 'constant', 'triSurface');
    await mkdir(triDir, { recursive: true });
    await copyFile(stlFile, path.join(triDir, 'enclosure.stl'));

    // 3. snappyHexMesh
    await exec('snappyHexMesh -overwrite', { cwd: caseDir, timeout: 300_000 });

    // 4. Check mesh quality
    const checkResult = await exec('checkMesh', { cwd: caseDir });
    return this.parseMeshCheck(checkResult.stdout);
  }

  async runSolver(caseDir: string, solver: OpenFOAMSolver): Promise<CFDResult> {
    const result = await exec(solver, {
      cwd: caseDir,
      timeout: 600_000,  // 10 min timeout
    });

    return {
      converged: this.checkConvergence(result.stdout),
      iterations: this.extractIterations(result.stdout),
      residuals: this.extractResiduals(result.stdout),
      ...await this.extractResults(caseDir),
    };
  }
}