diff --git a/pyproject.toml b/pyproject.toml index c302e6a..7e26f72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,5 +21,5 @@ dev = [ requires = ["setuptools>=61.0"] build-backend = "setuptools.build_meta" -[tool.uv] -python-version = "3.11" +[tool.setuptools] +py-modules = ["server"] diff --git a/server.py b/server.py index b2182b4..280c912 100644 --- a/server.py +++ b/server.py @@ -1,10 +1,13 @@ import os import json import logging +import mimetypes +import contextlib from typing import Any import httpx from fastmcp import FastMCP from starlette.applications import Starlette +from starlette.requests import Request from starlette.responses import JSONResponse from starlette.routing import Route, Mount from dotenv import load_dotenv @@ -491,6 +494,18 @@ Delete attachment. Params: id (required) ### attachments.redirect Get attachment URL. Params: id (required) +### Custom HTTP Gateway Upload +Rather than manually performing the multi-step attachments.create + PUT process, you can stream files directly through the MCP server's built-in HTTP streaming gateway: +- **Endpoint**: `POST /upload` +- **Query Parameters**: + - `documentId` (required): The UUID of the destination Outline document. + - `name` (required): The filename of the attachment. +- **Headers**: + - `Content-Length` (required): Total size of the file in bytes. + - `Content-Type` (optional): The MIME type of the file (auto-detected if omitted). +- **Body**: Raw binary stream of the file content. +- **Response**: JSON containing the created attachment's details, including its public access `url` on success. + --- ## Auth @@ -532,6 +547,167 @@ def get_api_reference() -> str: return API_REFERENCE +# --- Upload Endpoint --- +async def upload_endpoint(request: Request) -> JSONResponse: + """HTTP streaming endpoint to upload file attachments to Outline. + + This endpoint acts as a non-buffering gateway/proxy. It calls Outline's + attachments.create to register the metadata and obtain an upload URL, + then streams the incoming request body directly to S3. + """ + # 1. Authorization validation + auth_header = request.headers.get("Authorization") + if not auth_header or auth_header != f"Bearer {OUTLINE_API_TOKEN}": + return JSONResponse( + {"ok": False, "error": "Unauthorized"}, + status_code=401, + ) + + document_id = request.query_params.get("documentId") + if not document_id: + return JSONResponse( + {"ok": False, "error": "Missing 'documentId' query parameter"}, + status_code=400, + ) + + filename = request.query_params.get("name") + if not filename: + return JSONResponse( + {"ok": False, "error": "Missing 'name' query parameter"}, + status_code=400, + ) + + # Size is required by attachments.create + content_length_str = request.headers.get("content-length") + if not content_length_str: + return JSONResponse( + {"ok": False, "error": "Missing 'Content-Length' header"}, + status_code=400, + ) + + try: + size = int(content_length_str) + except ValueError: + return JSONResponse( + {"ok": False, "error": "Invalid 'Content-Length' header"}, + status_code=400, + ) + + # Validate file size limits (abuses and negative boundary checks) + MAX_UPLOAD_SIZE = int(os.getenv("MAX_UPLOAD_SIZE", "104857600")) # Default 100MB + if size < 0: + return JSONResponse( + {"ok": False, "error": "Invalid file size (must be non-negative)"}, + status_code=400, + ) + if size > MAX_UPLOAD_SIZE: + return JSONResponse( + {"ok": False, "error": f"File too large (exceeds limit of {MAX_UPLOAD_SIZE} bytes)"}, + status_code=413, + ) + + # Detect content type + content_type = ( + request.headers.get("content-type") + or mimetypes.guess_type(filename)[0] + or "application/octet-stream" + ) + + # 2. Initialize upload session in Outline + create_params = { + "name": filename, + "documentId": document_id, + "size": size, + "contentType": content_type, + } + + logger.info(f"Initiating attachment creation in Outline: {create_params}") + create_response = await outline_client.call("attachments.create", create_params) + + if not create_response.get("ok"): + return JSONResponse( + { + "ok": False, + "error": f"Failed to initialize attachment in Outline: {create_response.get('error')}", + }, + status_code=502, + ) + + attachment_data = create_response.get("data", {}) + upload_url = attachment_data.get("uploadUrl") + if not upload_url: + return JSONResponse( + {"ok": False, "error": "Outline API did not return an uploadUrl"}, + status_code=502, + ) + + # 3. Stream the request body directly to S3 upload_url + try: + # Pre-signed S3 PUT URLs usually enforce exact headers + headers = { + "Content-Type": content_type, + "Content-Length": str(size), + } + + logger.info(f"Streaming file content to storage (size={size} bytes)") + async with httpx.AsyncClient() as client: + put_response = await client.put( + upload_url, + content=request.stream(), + headers=headers, + timeout=600.0, + ) + put_response.raise_for_status() + + logger.info("Upload to storage completed successfully") + + # 4. Return success and attachment metadata + return JSONResponse( + { + "ok": True, + "data": { + "id": attachment_data.get("id"), + "name": attachment_data.get("name"), + "size": attachment_data.get("size"), + "contentType": attachment_data.get("contentType"), + "url": attachment_data.get("url"), + }, + } + ) + + except Exception as e: + # Clean up orphaned attachment record in Outline to maintain consistency + attachment_id = attachment_data.get("id") + if attachment_id: + logger.warning( + f"Cleaning up orphaned attachment {attachment_id} in Outline due to S3 upload failure" + ) + try: + await outline_client.call("attachments.delete", {"id": attachment_id}) + except Exception as cleanup_err: + logger.error( + f"Failed to delete orphaned attachment {attachment_id}: {str(cleanup_err)}" + ) + + if isinstance(e, httpx.HTTPStatusError): + logger.error( + f"Storage upload HTTP error: {e.response.status_code} - {e.response.text}" + ) + return JSONResponse( + { + "ok": False, + "error": f"Failed to upload attachment to storage: {str(e)}", + }, + status_code=502, + ) + else: + logger.error(f"Exception during attachment upload: {str(e)}") + return JSONResponse( + {"ok": False, "error": f"Upload stream failed: {str(e)}"}, + status_code=500, + ) + + # --- Health Check Endpoint --- async def health(request): """Health check endpoint for Docker/load balancer. @@ -554,16 +730,30 @@ async def health(request): # --- ASGI Application --- +@contextlib.asynccontextmanager +async def lifespan(app: Starlette): + """Custom ASGI lifespan manager that integrates with FastMCP's lifespan + and ensures the Outline HTTP client's connections are gracefully closed + on server shutdown. + """ + mcp_app = mcp.http_app() + async with mcp_app.lifespan(app): + yield + logger.info("Closing Outline client connections on server shutdown") + await outline_client.close() + + def create_app(): """Create the ASGI application with health check and MCP routes.""" mcp_app = mcp.http_app() routes = [ Route("/health", health, methods=["GET"]), + Route("/upload", upload_endpoint, methods=["POST"]), Mount("/", app=mcp_app), ] - return Starlette(routes=routes, lifespan=mcp_app.lifespan) + return Starlette(routes=routes, lifespan=lifespan) # Create the app instance for uvicorn diff --git a/test_upload.py b/test_upload.py new file mode 100644 index 0000000..286f139 --- /dev/null +++ b/test_upload.py @@ -0,0 +1,259 @@ +import json +from unittest.mock import AsyncMock, patch +import pytest +import httpx +from starlette.testclient import TestClient +from server import create_app, outline_client, OUTLINE_API_TOKEN + +# Helper mock responses +MOCK_ATTACH_CREATE_SUCCESS = { + "ok": True, + "data": { + "id": "attachment-5678-uuid", + "name": "test_stream.txt", + "size": 34, + "contentType": "text/plain", + "uploadUrl": "https://s3.example.com/fake-bucket/test_stream.txt?signature=xyz", + "url": "https://docs.example.com/attachments/attachment-5678-uuid" + } +} + +MOCK_ATTACH_CREATE_ERROR = { + "ok": False, + "error": "Document not found" +} + + +@pytest.mark.asyncio +async def test_upload_endpoint_success(): + """Verify that a successful stream upload coordinates with Outline, + streams content to S3, and returns the expected metadata. + """ + file_content = b"Hello, Outline Wiki Stream upload!" + content_length = len(file_content) + document_id = "doc-1234-uuid" + filename = "test_stream.txt" + content_type = "text/plain" + + app = create_app() + + # Mock OutlineClient attachments.create + with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call: + mock_call.return_value = MOCK_ATTACH_CREATE_SUCCESS + + # Mock S3 PUT response + mock_put_response = httpx.Response( + 200, request=httpx.Request("PUT", "https://s3.example.com") + ) + + # Mock httpx.AsyncClient to intercept put + with patch("httpx.AsyncClient") as MockClient: + mock_instance = AsyncMock() + mock_instance.put.return_value = mock_put_response + MockClient.return_value.__aenter__.return_value = mock_instance + + client = TestClient(app) + response = client.post( + f"/upload?documentId={document_id}&name={filename}", + content=file_content, + headers={ + "Content-Length": str(content_length), + "Content-Type": content_type, + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + + assert response.status_code == 200 + response_json = response.json() + assert response_json["ok"] is True + assert response_json["data"]["id"] == "attachment-5678-uuid" + + # Verify attachments.create parameters + mock_call.assert_called_once_with( + "attachments.create", + { + "name": filename, + "documentId": document_id, + "size": content_length, + "contentType": content_type + } + ) + + # Verify S3 PUT parameters + mock_instance.put.assert_called_once() + args, kwargs = mock_instance.put.call_args + assert args[0] == "https://s3.example.com/fake-bucket/test_stream.txt?signature=xyz" + assert kwargs["headers"]["Content-Type"] == content_type + assert kwargs["headers"]["Content-Length"] == str(content_length) + + +@pytest.mark.asyncio +async def test_upload_endpoint_unauthorized(): + """Verify that request with incorrect or missing auth is rejected with 401.""" + app = create_app() + client = TestClient(app) + + # Missing Authorization header + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + content=b"content", + headers={"Content-Length": "7"} + ) + assert response.status_code == 401 + assert response.json()["error"] == "Unauthorized" + + # Incorrect Authorization token + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + content=b"content", + headers={ + "Content-Length": "7", + "Authorization": "Bearer wrong-token" + } + ) + assert response.status_code == 401 + assert response.json()["error"] == "Unauthorized" + + +@pytest.mark.asyncio +async def test_upload_endpoint_oversized_file(): + """Verify that oversized files are rejected with 413 Payload Too Large.""" + app = create_app() + client = TestClient(app) + + # Content length greater than default MAX_UPLOAD_SIZE (100MB) + oversized_length = 104857600 + 1 # 100MB + 1 byte + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + headers={ + "Content-Length": str(oversized_length), + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + assert response.status_code == 413 + assert "File too large" in response.json()["error"] + + +@pytest.mark.asyncio +async def test_upload_endpoint_negative_size(): + """Verify that negative file size content lengths are rejected with 400.""" + app = create_app() + client = TestClient(app) + + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + headers={ + "Content-Length": "-10", + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + assert response.status_code == 400 + assert "Invalid file size" in response.json()["error"] + + +@pytest.mark.asyncio +async def test_upload_endpoint_missing_params(): + """Verify missing query params are handled with 400.""" + app = create_app() + client = TestClient(app) + + # Missing documentId + response = client.post( + "/upload?name=test.txt", + content=b"content", + headers={ + "Content-Length": "7", + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + assert response.status_code == 400 + assert "documentId" in response.json()["error"] + + # Missing name + response = client.post( + "/upload?documentId=doc-123", + content=b"content", + headers={ + "Content-Length": "7", + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + assert response.status_code == 400 + assert "name" in response.json()["error"] + + +@pytest.mark.asyncio +async def test_upload_endpoint_missing_content_length(): + """Verify that missing Content-Length header is handled correctly.""" + app = create_app() + client = TestClient(app) + + # We send a generator to prevent TestClient from auto-populating Content-Length + def dummy_generator(): + yield b"chunk" + + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + content=dummy_generator(), + headers={"Authorization": f"Bearer {OUTLINE_API_TOKEN}"} + ) + assert response.status_code == 400 + assert "Content-Length" in response.json()["error"] + + +@pytest.mark.asyncio +async def test_upload_endpoint_outline_create_failure(): + """Verify that Outline registration failures are handled correctly.""" + app = create_app() + + with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call: + mock_call.return_value = MOCK_ATTACH_CREATE_ERROR + + client = TestClient(app) + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + content=b"content", + headers={ + "Content-Length": "7", + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + + assert response.status_code == 502 + assert "Failed to initialize attachment" in response.json()["error"] + + +@pytest.mark.asyncio +async def test_upload_endpoint_s3_failure_cleanup(): + """Verify that if the S3 upload fails, attachments.delete is called on Outline to clean up.""" + app = create_app() + + with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call: + mock_call.return_value = MOCK_ATTACH_CREATE_SUCCESS + + # S3 PUT failure response + mock_put_response = httpx.Response( + 500, request=httpx.Request("PUT", "https://s3.example.com") + ) + + with patch("httpx.AsyncClient") as MockClient: + mock_instance = AsyncMock() + mock_instance.put.return_value = mock_put_response + MockClient.return_value.__aenter__.return_value = mock_instance + + client = TestClient(app) + response = client.post( + "/upload?documentId=doc-123&name=test.txt", + content=b"content", + headers={ + "Content-Length": "7", + "Authorization": f"Bearer {OUTLINE_API_TOKEN}" + } + ) + + # S3 returns 500 so our gateway should return 502 + assert response.status_code == 502 + assert "Failed to upload attachment to storage" in response.json()["error"] + + # Verify that attachments.delete was called to rollback the creation + mock_call.assert_any_call("attachments.delete", {"id": "attachment-5678-uuid"})