# NiChart DLMUSE MCP server

Hosts [`cbica/nichart_dlmuse`](https://hub.docker.com/r/cbica/nichart_dlmuse) (T1 MRI
skull-strip + MUSE ROI segmentation) behind an MCP server on a GPU EC2 instance, so
Claude Code can run segmentations remotely instead of every user needing a local GPU.

## Why it's shaped this way

MCP tool arguments are JSON. A `.nii.gz` is tens of MB of binary, and segmentation
takes ~1-2 minutes on GPU — too slow and too large for a single blocking tool call.
So:

- **File transfer happens outside the MCP protocol**, over a plain authenticated
  `POST /upload` endpoint. Only a small `upload_id` ever flows through an MCP tool
  argument.
- **Jobs are async**: `run_dlmuse_segmentation` enqueues and returns immediately;
  `get_job_status` polls; `get_job_result` fetches the CSV inline plus download
  links for the mask files.
- **One worker, serialized**: it's one shared GPU, so only one `docker run` executes
  at a time, queued behind an `asyncio.Queue`.
- **Bearer-token auth** per team member on every endpoint except `/healthz`.
- **Private instance, no public web tier**: the app binds only `127.0.0.1`, and
  the security group opens nothing but SSH (22). Users reach it by SSH
  port-forwarding their laptop into that loopback port with their private key —
  there is no domain and nothing internet-facing to attack.
- **TLS is a self-signed cert**, generated once on the instance, since a
  publicly-issued cert (Let's Encrypt, etc.) requires the instance to be
  reachable from the internet for domain validation, which this one isn't.
  Each user trusts that one cert on their laptop (see below).
- **Scans don't live forever**: uploads and job outputs are deleted after
  `RETENTION_HOURS` (default 24h).

## Architecture

```
Claude Code (laptop)                    SSH tunnel                  EC2 (private, SG: 22 only)
  │  ssh -i key.pem -L 8420:127.0.0.1:8420 user@instance ─────────────────►  │
  │                                                                          │
  │  1. curl https://127.0.0.1:8420/upload ─────(via tunnel)──────────►  MCP server :8420 (127.0.0.1, self-signed TLS)
  │  2. run_dlmuse_segmentation ────────────────(via tunnel)──────────►         │
  │  3. get_job_status (poll) ──────────────────(via tunnel)──────────►  asyncio job queue (1 worker)
  │  4. get_job_result ─────────────────────────(via tunnel)──────────►         │
                                                                   docker run --gpus all cbica/nichart_dlmuse
```

## Repo layout

```
server/app.py     MCP tools (run_dlmuse_segmentation, get_job_status, get_job_result)
                  + HTTP routes (/upload, /download/{job_id}/{filename}, /healthz)
server/jobs.py    job queue/worker, docker invocation, root-owned-output cleanup
server/auth.py    bearer-token ASGI middleware
server/config.py  env-driven settings
deploy/           EC2 provisioning script (installs Docker, GPU toolkit, self-signed cert, systemd unit)
```

## One-time EC2 setup

Requires an existing GPU EC2 instance (AWS Deep Learning AMI recommended — NVIDIA
driver and Docker are usually already installed).

1. Copy this repo onto the instance (`git clone` / `scp -r`).
2. `cp .env.example .env` and fill in at least `TOKENS` — one `name:token` pair per
   team member, comma-separated. Generate tokens with `openssl rand -hex 32`.
3. In the instance's security group: allow inbound `22` (SSH) only, restricted to
   your team's IPs or a bastion. Open nothing else — no `443`, no `8420`. The app
   binds `127.0.0.1` and is never reachable except through an SSH tunnel.
4. Run the provisioning script:
   ```
   sudo ./deploy/setup_ec2.sh
   ```
   It's idempotent — installs Docker/nvidia-container-toolkit only if missing,
   pulls the DLMUSE image, creates a dedicated `nichart-mcp` service user, deploys
   the code to `/opt/nichart-mcp`, generates a self-signed TLS cert (SAN =
   `127.0.0.1`/`localhost`), and installs the `nichart-mcp` systemd service.
5. Verify, from the instance itself:
   ```
   curl --cacert /opt/nichart-mcp/tls/server.crt https://127.0.0.1:8420/healthz
   ```
6. Copy `/opt/nichart-mcp/tls/server.crt` off the instance so you can hand it to
   each team member (e.g. `scp -i key.pem ec2-user@<instance-ip>:/opt/nichart-mcp/tls/server.crt .`).

To add or revoke a user later: edit `TOKENS` in `/opt/nichart-mcp/.env` on the
instance, then `sudo systemctl restart nichart-mcp`.

### Updating the code

Re-run `sudo ./deploy/setup_ec2.sh` from an updated checkout on the instance — it
re-syncs `/opt/nichart-mcp` (leaving the existing TLS cert alone), reinstalls
dependencies, and restarts the service.

## Connecting from a laptop

Each team member needs: their SSH private key (for a Windows PuTTY `.ppk` key,
convert it once with `puttygen key.ppk -O private-openssh -o key.pem` so the
standard `ssh` client can use it), their bearer token, and the `server.crt` file
from setup step 6.

**1. Trust the self-signed cert once**, so `curl`/Claude Code stop rejecting it:
- macOS: `security add-trusted-cert -d -r trustRoot -k ~/Library/Keychains/login.keychain-db server.crt`
- Linux: `sudo cp server.crt /usr/local/share/ca-certificates/nichart-mcp.crt && sudo update-ca-certificates`
- Windows: `certutil -addstore -f "ROOT" server.crt`

**2. Open the SSH tunnel** (leave this running in a terminal while using Claude Code):
```
ssh -i key.pem -N -L 8420:127.0.0.1:8420 <ssh_user>@<instance-ip>
```
If the instance has no public IP and you only reach it through a bastion, add
`-J <bastion_user>@<bastion_host>`.

**3. Register the MCP server** with Claude Code, once:
```
claude mcp add --transport http nichart-dlmuse https://127.0.0.1:8420/mcp \
  --header "Authorization: Bearer <their-token>"
```

## Using it

With the tunnel open, ask Claude Code to segment a scan; it will:

1. Upload the file (through the tunnel, so `127.0.0.1:8420` is correct even
   though the file is headed for the remote instance):
   ```
   curl -X POST -H "Authorization: Bearer <token>" \
     -F "file=@/path/to/scan.nii.gz" \
     https://127.0.0.1:8420/upload
   # -> {"upload_id": "..."}
   ```
2. Call the `run_dlmuse_segmentation` tool with that `upload_id` -> gets a `job_id`.
3. Poll `get_job_status(job_id)` until `status == "done"` (typically ~1-2 min on
   GPU).
4. Call `get_job_result(job_id)` -> ROI volumes CSV inline, plus
   `/download/{job_id}/{filename}` links for the ICV and MUSE mask NIfTI files
   (fetch those from `https://127.0.0.1:8420/download/...` with the same bearer
   token, through the same tunnel).

## Operating notes

- **GPU concurrency**: one job runs at a time by design (single shared GPU). A
  busy team will see jobs queue; `get_job_status` reports `queue_position`.
- **Job state is in-memory**: a `systemctl restart nichart-mcp` loses in-flight
  job records (the uploaded scan and any partial output on disk are unaffected,
  but you'd need to resubmit). Fine at small-team scale; if you outgrow it,
  swap the in-memory dict in `server/jobs.py` for Redis/RQ.
- **PHI**: scans are real patient data. `RETENTION_HOURS` bounds how long they sit
  on disk, but confirm that's compatible with your data-handling requirements before
  pointing this at real patients. Consider enabling EBS encryption on the instance's
  volume if you haven't already.
- **The DLMUSE container runs as root** inside (it hardcodes writes to
  `/app/pipeline.log`, so it can't run under `--user`). Its output ends up
  root-owned; `server/jobs.py`'s cleanup falls back to a throwaway `alpine`
  container to force-remove those directories.

## Local development

```
python3 -m venv .venv && .venv/bin/pip install -r requirements.txt
cp .env.example .env   # fill in TOKENS
TOKENS=dev:devtoken DATA_DIR=/tmp/nichart-dev .venv/bin/python -m server.app
```

This runs the full server (auth, upload, job queue, MCP tools) locally. Actually
running a segmentation still requires Docker with GPU access — on a machine
without a GPU, jobs will fail at the `docker run` step but everything else
(routing, auth, queueing, status/error reporting) is exercisable.