Module 6: Project — Advanced Task Manager API

WebSocket, File Uploads, and Middleware

Overview

In capsule 02 you registered ws_router and uploads_router in main.py, but the files that export them don't exist yet — that's why the app still doesn't start. In this capsule you implement them: the WebSocket endpoint that exposes the ConnectionManager you already built, the uploads router with validation, and you verify that the logging+timing middleware (also from capsule 02) is measuring every request correctly.

By the end of this capsule, the API runs without errors and the 3 advanced features — real-time, files, and observability — work end-to-end.


WebSocket Router

Create app/routers/websocket.py:

import logging
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.websocket_manager import ws_manager
from app.data import tasks_db

logger = logging.getLogger("task_manager.websocket")

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await ws_manager.connect(websocket)
    logger.info(f"Client connected | Total: {ws_manager.connection_count}")

    await ws_manager.send_personal({
        "type": "connection_established",
        "data": {
            "message": "Connected to the Task Manager stream",
            "active_tasks": len(tasks_db),
            "your_position": ws_manager.connection_count
        }
    }, websocket)

    try:
        while True:
            message = await websocket.receive_json()
            action = message.get("action")

            if action == "ping":
                await ws_manager.send_personal(
                    {"type": "pong", "data": {"echo": message.get("data", {})}},
                    websocket
                )

            elif action == "status":
                await ws_manager.send_personal({
                    "type": "status_response",
                    "data": {
                        "connections": ws_manager.connection_count,
                        "tasks_count": len(tasks_db)
                    }
                }, websocket)

            elif action == "broadcast_test":
                await ws_manager.broadcast_event("test_broadcast", {
                    "from": "client",
                    "message": message.get("data", {}).get("message", "ping")
                })

            else:
                await ws_manager.send_personal({
                    "type": "error",
                    "data": {
                        "error": "Unknown action",
                        "received": action,
                        "valid_actions": ["ping", "status", "broadcast_test"]
                    }
                }, websocket)

    except WebSocketDisconnect:
        ws_manager.disconnect(websocket)
        logger.info(f"Client disconnected | Total: {ws_manager.connection_count}")

Design decisions

No prefix on the router — WebSocket endpoints usually live at the root (/ws) with no namespace. If you want to version it, use APIRouter(prefix="/v1").

A welcome message with send_personal — It confirms to the client that the connection works and gives it some initial context (how many tasks there are, its position).

Client actions: ping, status, broadcast_test — The protocol is JSON with an action field. Keep the set small and well documented.

broadcast_test — Useful during development to check that the broadcast reaches every connected client. Remove it in production.

WebSocketDisconnect closes cleanly — Without try/except, the manager holds on to dead connections that will fail on the next broadcast.


Uploads Router

Create app/routers/uploads.py:

import os
import secrets
import logging
from pathlib import Path
from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Depends, BackgroundTasks
from fastapi.responses import FileResponse
from app.models import AttachmentInfo
from app.dependencies import get_task_or_404
from app.websocket_manager import ws_manager
from app.background.notifications import log_audit_event

logger = logging.getLogger("task_manager.uploads")

router = APIRouter(prefix="/tasks", tags=["Uploads"])

UPLOAD_DIR = Path("app/uploads")
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

ALLOWED_CONTENT_TYPES = {
    "image/jpeg",
    "image/png",
    "image/gif",
    "application/pdf",
    "text/plain",
    "text/csv",
    "application/json"
}

MAX_FILE_SIZE_MB = 5
MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024


def validate_upload(file: UploadFile, size_bytes: int) -> None:
    if file.content_type not in ALLOWED_CONTENT_TYPES:
        raise HTTPException(
            status_code=415,
            detail={
                "error": "File type not allowed",
                "received": file.content_type,
                "allowed": sorted(ALLOWED_CONTENT_TYPES)
            }
        )

    if size_bytes > MAX_FILE_SIZE_BYTES:
        raise HTTPException(
            status_code=413,
            detail={
                "error": "File too large",
                "max_size_mb": MAX_FILE_SIZE_MB,
                "received_mb": round(size_bytes / (1024 * 1024), 2)
            }
        )

    if size_bytes == 0:
        raise HTTPException(
            status_code=400,
            detail={"error": "Empty file"}
        )


def safe_filename(original: str) -> str:
    extension = Path(original).suffix.lower()
    token = secrets.token_hex(8)
    return f"{token}{extension}"


@router.post(
    "/{task_id}/attachments",
    status_code=201,
    response_model=AttachmentInfo
)
async def upload_attachment(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    description: str = Form(default=""),
    task: dict = Depends(get_task_or_404)
):
    contents = await file.read()
    size_bytes = len(contents)

    validate_upload(file, size_bytes)

    stored_name = safe_filename(file.filename or "upload.bin")
    file_path = UPLOAD_DIR / stored_name

    with file_path.open("wb") as buffer:
        buffer.write(contents)

    attachment = {
        "filename": stored_name,
        "original_name": file.filename or "upload.bin",
        "description": description,
        "content_type": file.content_type,
        "size_bytes": size_bytes
    }
    task.setdefault("attachments", []).append(attachment)

    background_tasks.add_task(
        log_audit_event,
        action="upload_attachment",
        details={
            "task_id": task["id"],
            "filename": stored_name,
            "size_bytes": size_bytes
        }
    )

    await ws_manager.broadcast_event("attachment_added", {
        "task_id": task["id"],
        "filename": stored_name,
        "original_name": attachment["original_name"]
    })

    logger.info(
        f"Upload OK | task={task['id']} | "
        f"file={attachment['original_name']}{stored_name} | {size_bytes} bytes"
    )

    return attachment


@router.get(
    "/{task_id}/attachments",
    response_model=list[AttachmentInfo]
)
async def list_attachments(task: dict = Depends(get_task_or_404)):
    return task.get("attachments", [])


@router.get("/{task_id}/attachments/{filename}")
async def download_attachment(
    filename: str,
    task: dict = Depends(get_task_or_404)
):
    if "/" in filename or "\\" in filename or ".." in filename:
        raise HTTPException(status_code=400, detail="Invalid filename")

    attachment = next(
        (a for a in task.get("attachments", []) if a["filename"] == filename),
        None
    )
    if not attachment:
        raise HTTPException(
            status_code=404,
            detail={
                "error": "Attachment not found on this task",
                "task_id": task["id"],
                "filename": filename
            }
        )

    file_path = UPLOAD_DIR / filename
    if not file_path.exists():
        raise HTTPException(
            status_code=410,
            detail="The file was removed from storage"
        )

    return FileResponse(
        path=file_path,
        media_type=attachment["content_type"],
        filename=attachment["original_name"]
    )


@router.delete("/{task_id}/attachments/{filename}", status_code=204)
async def delete_attachment(
    filename: str,
    background_tasks: BackgroundTasks,
    task: dict = Depends(get_task_or_404)
):
    attachments = task.get("attachments", [])
    attachment = next((a for a in attachments if a["filename"] == filename), None)

    if not attachment:
        raise HTTPException(status_code=404, detail="Attachment not found")

    file_path = UPLOAD_DIR / filename
    if file_path.exists():
        file_path.unlink()

    attachments.remove(attachment)

    background_tasks.add_task(
        log_audit_event,
        action="delete_attachment",
        details={"task_id": task["id"], "filename": filename}
    )

Key points

UPLOAD_DIR.mkdir(parents=True, exist_ok=True) — Guarantees the folder exists when the module is imported. It doesn't break if it already exists.

secrets.token_hex(8) for names — 16 random hex characters. It avoids collisions between files with the same original name (e.g. two screenshot.png) and prevents path traversal from using the user's name directly.

Validation in order: type → size → empty — Fail fast and with specific messages. The client knows exactly what to fix.

UploadFile + Form in the same endpoint — To send file and description together, the client has to use multipart/form-data. FastAPI separates the fields automatically.

Path traversal blocked on download — Even though secrets.token_hex avoids malicious names on upload, we validate /, \, and .. on download for defense in depth.

410 Gone for removed files — If the metadata exists but the file doesn't, we return 410 (not 404). It signals that it existed but no longer does.

A WS broadcast on upload, but not on download — Only events that change state get broadcast. Downloads are pure reads.


About the middleware

The logging+timing middleware is already implemented in app/middleware/logging.py since capsule 02:

class LoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        start_time = time.time()
        response = await call_next(request)
        duration_ms = (time.time() - start_time) * 1000

        logger.info(
            f"{request.method} {request.url.path} "
            f"→ {response.status_code} ({duration_ms:.1f}ms)"
        )
        response.headers["X-Process-Time-Ms"] = f"{duration_ms:.1f}"
        return response

And it gets registered in main.py:

app.add_middleware(LoggingMiddleware)
app.add_middleware(CORSMiddleware, allow_origins=["*"], ...)

A quick middleware check:

curl -i http://localhost:8000/health | grep -i "x-process-time"
# X-Process-Time-Ms: 1.3

If you see the header, the middleware is working on every request. If it doesn't show up, check that app.add_middleware(LoggingMiddleware) comes before app.include_router(...) in main.py.

A note on WebSocket and HTTP middleware: Middleware based on BaseHTTPMiddleware does not run for WebSocket connections (they're a different protocol). For WS logging, use the logger.info() calls inside the WS endpoint itself.


Partial verification

1. Start the server

uvicorn app.main:app --reload

You should see this in the logs:

INFO [task_manager] Task Manager API started with sample data
INFO:     Application startup complete.

If it fails with ImportError: cannot import name 'router' from 'app.routers.websocket', you're missing the file from this capsule.

2. Test uploads with curl

# Create a test file
echo "Project notes" > /tmp/notes.txt

# Upload
curl -s -X POST http://localhost:8000/tasks/1/attachments \
  -F "file=@/tmp/notes.txt" \
  -F "description=Initial notes" | python -m json.tool

# List the task's attachments
curl -s http://localhost:8000/tasks/1/attachments | python -m json.tool

# Download (replace FILENAME with the real filename from the previous response)
curl -s "http://localhost:8000/tasks/1/attachments/FILENAME.txt" -o /tmp/downloaded.txt
diff /tmp/notes.txt /tmp/downloaded.txt && echo "Files are identical"

3. Test the validations

# Type not allowed
echo '<?xml version="1.0"?>' > /tmp/test.xml
curl -s -X POST http://localhost:8000/tasks/1/attachments \
  -F "file=@/tmp/test.xml;type=application/xml" | python -m json.tool
# Expected: 415 with the list of allowed types

# File too large (generates 6MB of zeros)
dd if=/dev/zero of=/tmp/big.pdf bs=1M count=6 2>/dev/null
curl -s -X POST http://localhost:8000/tasks/1/attachments \
  -F "file=@/tmp/big.pdf;type=application/pdf" | python -m json.tool
# Expected: 413 with max_size_mb

4. Test the WebSocket

Install websocat (a CLI client for WS):

# macOS
brew install websocat

# Linux
cargo install websocat
# or download the binary from github.com/vi/websocat/releases

Connect and try it:

websocat ws://localhost:8000/ws

The message you should get on connecting:

{
  "type": "connection_established",
  "data": {"message": "Connected to the Task Manager stream", ...}
}

Send actions (type them into the websocat terminal):

{"action": "ping", "data": {"hello": "world"}}
{"action": "status"}

And while websocat is connected, in another terminal:

curl -s -X POST http://localhost:8000/tasks/ \
  -H "Content-Type: application/json" \
  -d '{"title": "Test WS", "priority": "high"}'

In websocat you'll see the task_created event arrive. That confirms the whole flow: HTTP request → background tasks → WS broadcast.

5. Check the timing in the headers

curl -s -i http://localhost:8000/tasks/ | head -20

Look for X-Process-Time-Ms: 2.4 (or similar). It confirms the middleware is measuring.


Troubleshooting

RuntimeError: Form data requires "python-multipart"

The dependency is missing. Install it with pip install python-multipart and restart the server.

The upload returns 422 without touching your validation

The 422 comes from Pydantic/FastAPI before your code. Check that the client is sending multipart/form-data, not application/json. With curl, use -F, not -d.

The WebSocket closes immediately with code 1006

The server most likely raised an exception inside the endpoint before accept(). Check the server logs — it's usually a missing import or a misspelled name.

WS broadcasts don't reach connected clients

Confirm that the endpoint calls await ws_manager.connect(websocket) (not just accept()). If you use accept() directly, the client never gets added to active_connections and the broadcasts skip it.

FileResponse returns 200 but the file is empty

Check that UPLOAD_DIR resolves to the same absolute path for upload and download. If you start uvicorn from a folder other than the project's, the files get written somewhere else.

Path traversal with .. in download

The if "/" in filename or "\\" in filename or ".." in filename validation rejects it. If an edge case reaches you (e.g. filename = "%2e%2e/etc/passwd"), FastAPI already decodes it before it gets to your handler — the check still works because it sees the real characters.


The project structure at the end of this capsule

task-manager-api/
├── app/
│   ├── __init__.py
│   ├── main.py                 ✅ cap. 02
│   ├── models.py               ✅ cap. 02
│   ├── data.py                 ✅ cap. 02
│   ├── dependencies.py         ✅ cap. 02
│   ├── websocket_manager.py    ✅ cap. 02
│   ├── routers/
│   │   ├── __init__.py         ✅ cap. 02
│   │   ├── tasks.py            ✅ cap. 03
│   │   ├── uploads.py          ✅ cap. 04 ← this capsule
│   │   └── websocket.py        ✅ cap. 04 ← this capsule
│   ├── middleware/
│   │   ├── __init__.py         ✅ cap. 02
│   │   └── logging.py          ✅ cap. 02
│   ├── background/
│   │   ├── __init__.py         ✅ cap. 02
│   │   └── notifications.py    ✅ cap. 02
│   └── uploads/                (auto-created by uploads.py)
└── requirements.txt            ⏳ cap. 05
└── README.md                   ⏳ cap. 05

At this point you have a working API with all the features. All that's left is packaging it with setup, a README, and an end-to-end verification.


Resources

  1. FastAPI WebSockets — The official WebSocket patterns
  2. FastAPI Request FilesUploadFile, File, Form
  3. Starlette FileResponse — Download details
  4. websocat — A CLI client for testing WebSockets
  5. OWASP File Upload — Common vulnerabilities and mitigations
  6. Python secrets module — Generating cryptographically secure tokens

What's next?

In Capsule 05 you verify that everything works end-to-end with a test script, write the professional README.md, add tests with TestClient, and learn how to present this project in your portfolio. It's the close of the guide.