# MLX MCP Server 🚀

**Model Context Protocol server for ML Offload System**

Unified interface for LLM inference management via [Model Context Protocol (MCP)](https://modelcontextprotocol.io).
Control your ML models directly from **Claude Desktop**, **VSCode**, **Zed Editor**, or **Amazon Bedrock**.

---

## Features

### Core Functionality
- 🧠 **Model Registry Access** - List, search, and inspect ML models
- 🎮 **Backend Control** - Load/unload/switch models across backends (Ollama, llama.cpp, vLLM, TGI)
- 📊 **VRAM Monitoring** - Real-time GPU memory status
- ⚡ **Hot-Reload** - Switch models without restarting services
- 🔧 **Auto-Discovery** - Trigger model scans to update registry
- 🏥 **Health Checks** - Monitor ML Offload API status

### 💰 Token Economy (NEW!)
- **Smart Caching** - Per-tool TTL (5min for models, 10s for VRAM)
- **Auto-Summarization** - Concise responses by default (~80-90% token savings)
- **Rate Limiting** - Prevents API abuse and excessive costs
- **Cache Statistics** - Monitor hit rate and savings
- **Verbose Mode** - Full JSON when needed (`verbose: true`)

**Cost Savings**: ~80% reduction in API costs (Amazon/Anthropic billing)

---

## Prerequisites

1. **ML Offload API** running on `http://localhost:9000`
   (See `/etc/nixos/modules/ml/offload` for setup)

2. **Node.js 18+**
   ```bash
   node --version  # Should be 18.0.0 or higher
   ```

---

## Installation

```bash
# Clone or navigate to this directory
cd ~/dev/mlx-mcp

# Install dependencies
npm install

# Build TypeScript
npm run build
```

---

## Configuration

### Environment Variables

```bash
# Optional: Change ML Offload API URL (default: http://localhost:9000)
export ML_OFFLOAD_API_URL="http://localhost:9000"
```

---

## Usage

### 1. **Claude Desktop** (Recommended)

Add to your Claude Desktop config (`~/Library/Application Support/Claude/claude_desktop_config.json` on macOS):

```json
{
  "mcpServers": {
    "mlx": {
      "command": "node",
      "args": ["/home/kernelcore/dev/mlx-mcp/dist/index.js"],
      "env": {
        "ML_OFFLOAD_API_URL": "http://localhost:9000"
      }
    }
  }
}
```

**Restart Claude Desktop** and you'll see MLX tools available!

---

### 2. **VSCode / VSCodium** (via Continue extension)

1. Install [Continue extension](https://marketplace.visualstudio.com/items?itemName=Continue.continue)
2. Add to `.continue/config.json`:

```json
{
  "experimental": {
    "modelContextProtocol": true
  },
  "mcpServers": {
    "mlx": {
      "command": "node",
      "args": ["/home/kernelcore/dev/mlx-mcp/dist/index.js"]
    }
  }
}
```

3. Reload VSCode and access via Continue chat

---

### 3. **Zed Editor**

Add to Zed settings (`~/.config/zed/settings.json`):

```json
{
  "context_servers": {
    "mlx-mcp": {
      "command": "node",
      "args": ["/home/kernelcore/dev/mlx-mcp/dist/index.js"]
    }
  }
}
```

---

### 4. **Amazon Bedrock**

Use with Bedrock Agent runtime via stdio:

```python
import subprocess
import json

# Start MCP server
process = subprocess.Popen(
    ["node", "/home/kernelcore/dev/mlx-mcp/dist/index.js"],
    stdin=subprocess.PIPE,
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE
)

# Send MCP request
request = {
    "jsonrpc": "2.0",
    "id": 1,
    "method": "tools/call",
    "params": {
        "name": "list_models",
        "arguments": {}
    }
}

process.stdin.write((json.dumps(request) + "\n").encode())
process.stdin.flush()

# Read response
response = process.stdout.readline()
print(json.loads(response))
```

---

## Available Tools

### 📋 Model Management

- **`list_models`** - List all models in registry
  - Filters: `format`, `backend`, `limit`
  - Example: "List all GGUF models compatible with llamacpp"

- **`get_model_info`** - Get detailed model information
  - Args: `model_id`
  - Example: "Show details for model ID 5"

- **`trigger_model_scan`** - Scan `/var/lib/ml-models` for new models
  - No args required
  - Example: "Scan for new models"

### ⚙️ Backend Operations

- **`list_backends`** - Show all available backends
  - No args required
  - Example: "Show available backends"

- **`load_model`** - Load model on backend
  - Args: `model_id`, `backend`, `priority` (optional), `gpu_layers` (optional)
  - Example: "Load model 3 on llamacpp with 24 GPU layers"

- **`unload_model`** - Unload model from backend
  - Args: `backend`
  - Example: "Unload model from ollama"

- **`switch_model`** - Hot-switch model on backend
  - Args: `backend`, `model_id`, `gpu_layers` (optional)
  - Example: "Switch llamacpp to model 7"

### 📊 Monitoring

- **`get_vram_status`** - Real-time GPU VRAM status
  - No args required
  - Example: "Show VRAM usage"

- **`get_status`** - Complete system status
  - Shows VRAM, backends, loaded models
  - Example: "Show system status"

- **`health_check`** - Check ML Offload API health
  - No args required
  - Example: "Is the ML API healthy?"

---

## Example Prompts (for Claude/AI Assistants)

```
"List all GGUF models smaller than 5GB"
→ Calls: list_models with format filter

"What's the VRAM usage right now?"
→ Calls: get_vram_status

"Load mistral-7b-q4 on llamacpp with 28 GPU layers"
→ Calls: load_model with specific params

"Show me model 12's details"
→ Calls: get_model_info with model_id=12

"Scan for new models I just added"
→ Calls: trigger_model_scan

"Switch ollama to the new qwen model (ID 18)"
→ Calls: switch_model with backend=ollama, model_id=18
```

---

## Development

```bash
# Watch mode (auto-rebuild on changes)
npm run dev

# Build
npm run build

# Run
npm start
```

---

## Troubleshooting

### "Connection refused" error

**Problem**: MCP server can't reach ML Offload API

**Solution**:
```bash
# Check if ML Offload API is running
curl http://localhost:9000/health

# If not, enable it in NixOS config:
# /etc/nixos/hosts/kernelcore/configuration.nix
kernelcore.ml.offload.enable = true;
kernelcore.ml.offload.api.enable = true;

# Rebuild NixOS
sudo nixos-rebuild switch
```

### Tools not appearing in Claude Desktop

**Problem**: MCP tools not showing up

**Solution**:
1. Check config file path is correct
2. Verify `dist/index.js` exists (run `npm run build`)
3. Restart Claude Desktop completely
4. Check logs in Console.app (macOS) or journalctl (Linux)

### "Permission denied" error

**Problem**: Can't execute server

**Solution**:
```bash
chmod +x /home/kernelcore/dev/mlx-mcp/dist/index.js
```

---

## Architecture

```
┌─────────────────────────────────────────────────────────────┐
│                    AI Assistant (Claude/VSCode/Zed)         │
└────────────────────────┬────────────────────────────────────┘
                         │ MCP Protocol (stdio)
                         │
┌────────────────────────▼────────────────────────────────────┐
│                     MLX MCP Server (TypeScript)             │
│  - Tool handlers                                            │
│  - Request/Response formatting                              │
│  - Error handling                                           │
└────────────────────────┬────────────────────────────────────┘
                         │ HTTP REST API
                         │
┌────────────────────────▼────────────────────────────────────┐
│              ML Offload Manager API (Rust + Axum)           │
│  - Model Registry (SQLite)                                  │
│  - VRAM Intelligence (nvidia-smi)                           │
│  - Backend Orchestration                                    │
└────────────────────────┬────────────────────────────────────┘
                         │
        ┌────────────────┼────────────────┐
        │                │                │
        ▼                ▼                ▼
    ┌────────┐      ┌─────────┐     ┌──────┐
    │ Ollama │      │llama.cpp│     │ vLLM │
    └────────┘      └─────────┘     └──────┘
```

---

## License

MIT - See LICENSE file

---

## Support

- ML Offload Documentation: `/etc/nixos/docs/`
- MCP Protocol Docs: https://modelcontextprotocol.io
- Issues: File in this repository

---

**Made with ❤️ by kernelcore**