mcp-upload-artifact/server.py

"""
upload-artifact MCP server — upload files/text/images to the chat-artifacts S3 bucket.

The chat-artifacts bucket on Garage S3 is public-read (website endpoint).
Objects are accessible at: https://chat-artifacts.s3.nuclide.systems/<key>

Tools:
  upload_text      — upload a text string (Markdown, HTML, SVG, CSV, …)
  upload_base64    — upload binary content supplied as base64
  list_artifacts   — list recent artifacts (optionally filtered by prefix)
  get_url          — get the public URL for an artifact key
  delete_artifact  — delete an artifact by key
"""
from __future__ import annotations

import base64
import mimetypes
import os
from datetime import datetime, timezone

import boto3
from botocore.config import Config
from fastmcp import FastMCP

ENDPOINT    = os.environ.get("CHAT_ARTIFACTS_ENDPOINT", "https://s3.nuclide.systems")
ACCESS_KEY  = os.environ.get("CHAT_ARTIFACTS_ACCESS_KEY", "")
SECRET_KEY  = os.environ.get("CHAT_ARTIFACTS_SECRET_KEY", "")
BUCKET      = os.environ.get("CHAT_ARTIFACTS_BUCKET", "chat-artifacts")
PUBLIC_BASE = "https://chat-artifacts.s3.nuclide.systems"

mcp = FastMCP("upload-artifact")


def _s3():
    return boto3.client(
        "s3",
        endpoint_url=ENDPOINT,
        aws_access_key_id=ACCESS_KEY,
        aws_secret_access_key=SECRET_KEY,
        region_name="garage",
        config=Config(
            signature_version="s3v4",
            request_checksum_calculation="when_required",
            response_checksum_validation="when_required",
        ),
    )


def _date_prefix() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d/")


@mcp.tool()
def upload_text(
    content: str,
    filename: str,
    content_type: str = "text/plain; charset=utf-8",
    prefix: str = "",
) -> dict:
    """Upload a text string to the chat-artifacts bucket and return a public URL.

    content_type examples: 'text/markdown', 'text/html', 'image/svg+xml', 'text/csv'.
    prefix: optional subfolder (e.g. 'plots/', 'reports/'). Auto-date-prefixed if empty.
    """
    key = f"{prefix or _date_prefix()}{filename}"
    body = content.encode("utf-8")
    _s3().put_object(Bucket=BUCKET, Key=key, Body=body, ContentType=content_type)
    return {"key": key, "url": f"{PUBLIC_BASE}/{key}", "bytes": len(body)}


@mcp.tool()
def upload_base64(
    data_b64: str,
    filename: str,
    content_type: str = "application/octet-stream",
    prefix: str = "",
) -> dict:
    """Upload binary content (base64-encoded) to the chat-artifacts bucket.

    Useful for PNG/JPEG images, PDFs, ZIP archives.
    data_b64: standard or URL-safe base64; padding is optional.
    Returns the public URL.
    """
    data = base64.b64decode(data_b64 + "==")
    key = f"{prefix or _date_prefix()}{filename}"
    if content_type == "application/octet-stream":
        guessed, _ = mimetypes.guess_type(filename)
        if guessed:
            content_type = guessed
    _s3().put_object(Bucket=BUCKET, Key=key, Body=data, ContentType=content_type)
    return {"key": key, "url": f"{PUBLIC_BASE}/{key}", "bytes": len(data)}


@mcp.tool()
def list_artifacts(prefix: str = "", limit: int = 20) -> list[dict]:
    """List artifacts in the chat-artifacts bucket, newest first.

    prefix: filter by key prefix (e.g. '2026-05-21/' or 'plots/').
    Returns list of {key, bytes, last_modified, url}.
    """
    resp = _s3().list_objects_v2(Bucket=BUCKET, Prefix=prefix, MaxKeys=limit)
    objs = sorted(resp.get("Contents", []), key=lambda o: o["LastModified"], reverse=True)
    return [
        {
            "key": o["Key"],
            "bytes": o["Size"],
            "last_modified": o["LastModified"].isoformat(),
            "url": f"{PUBLIC_BASE}/{o['Key']}",
        }
        for o in objs
    ]


@mcp.tool()
def get_url(key: str) -> str:
    """Return the public URL for a known artifact key."""
    return f"{PUBLIC_BASE}/{key}"


@mcp.tool()
def delete_artifact(key: str) -> str:
    """Delete an artifact by key. This is irreversible — confirm with the user first."""
    _s3().delete_object(Bucket=BUCKET, Key=key)
    return f"Deleted: {key}"


if __name__ == "__main__":
    import asyncio
    asyncio.run(mcp.run_http_async(host="0.0.0.0", port=8000))