From 487f5355a0a6ce5a1ef80e11b46ae8629a32a2a3 Mon Sep 17 00:00:00 2001
From: Ben <b3nw@duck.com>
Date: Fri, 2 Jan 2026 02:49:43 +0000
Subject: [PATCH] Fix UniFi OS authentication and simplify server architecture

- Use /api/auth/login for UniFi OS controllers (UDM, Cloud Gateway)
- Use /api/login for standalone controllers
- Refactor to use FastMCP's native mcp.run() with custom_route for /health
- Switch to network_mode: host for local network access
- Include README.md in Dockerfile for hatchling build
---
 Dockerfile         |  4 ++--
 docker-compose.yml |  3 +--
 server.py          | 45 ++++++++++-----------------------------------
 unifi_client.py    | 13 +++++++++++--
 4 files changed, 24 insertions(+), 41 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index fe2f86a..2d79aed 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -6,8 +6,8 @@ WORKDIR /app
 # Install build dependencies
 RUN pip install --no-cache-dir uv
 
-# Copy dependency files
-COPY pyproject.toml ./
+# Copy dependency files and README (required by hatchling)
+COPY pyproject.toml README.md ./
 
 # Create virtual environment and install dependencies
 RUN uv venv /app/.venv
diff --git a/docker-compose.yml b/docker-compose.yml
index 4a1f675..03db079 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -4,8 +4,7 @@ services:
     image: gitea.ext.ben.io/b3nw/unifi-mcp-light:latest
     container_name: unifi-mcp-light
     restart: unless-stopped
-    ports:
-      - "${PORT:-8100}:8000"
+    network_mode: host
     environment:
       - UNIFI_HOST=${UNIFI_HOST}
       - UNIFI_USERNAME=${UNIFI_USERNAME}
diff --git a/server.py b/server.py
index f8c8ee6..11ce4c2 100644
--- a/server.py
+++ b/server.py
@@ -11,18 +11,14 @@ Following the BLUEPRINT.md pattern:
 - Starlette wrapper with health endpoint
 """
 
-import asyncio
 import json
 import logging
 import os
-from contextlib import asynccontextmanager
-from typing import Any, Optional
+from typing import Optional
 
 from dotenv import load_dotenv
 from fastmcp import FastMCP
-from starlette.applications import Starlette
 from starlette.responses import JSONResponse
-from starlette.routing import Mount, Route
 
 from api_docs import get_api_docs
 from unifi_client import UnifiClient, UnifiClientError, UnifiWriteBlockedError
@@ -355,11 +351,14 @@ async def unifi_api_call(
 
 
 # =============================================================================
-# Starlette Wrapper with Health Endpoint
+# Health Check Endpoint (using FastMCP custom_route)
 # =============================================================================
 
+from starlette.requests import Request
 
-async def health_endpoint(request) -> JSONResponse:
+
+@mcp.custom_route("/health", methods=["GET"])
+async def health_endpoint(request: Request) -> JSONResponse:
     """
     Health check endpoint for Docker/Kubernetes.
 
@@ -386,9 +385,8 @@ async def health_endpoint(request) -> JSONResponse:
         )
 
 
-@asynccontextmanager
-async def lifespan(app):
-    """Application lifespan manager."""
+def main():
+    """Entry point for the server."""
     logger.info("Starting UniFi MCP Light server...")
     logger.info(f"Controller: {CONFIG['host']}:{CONFIG['port']}")
     logger.info(f"Site: {CONFIG['site']}")
@@ -396,31 +394,8 @@ async def lifespan(app):
         f"Write operations: {'enabled' if CONFIG['allow_writes'] else 'disabled'}"
     )
 
-    yield
-
-    logger.info("Shutting down UniFi MCP Light server...")
-    await close_client()
-
-
-def create_app() -> Starlette:
-    """Create the Starlette ASGI application."""
-    mcp_app = mcp.http_app()
-
-    return Starlette(
-        routes=[
-            Route("/health", health_endpoint),
-            Mount("/", app=mcp_app),
-        ],
-        lifespan=lifespan,
-    )
-
-
-def main():
-    """Entry point for the server."""
-    import uvicorn
-
-    app = create_app()
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    # Run with HTTP transport - MCP endpoint at /mcp, health at /health
+    mcp.run(transport="http", host="0.0.0.0", port=8000)
 
 
 if __name__ == "__main__":
diff --git a/unifi_client.py b/unifi_client.py
index ca79545..693b1f1 100644
--- a/unifi_client.py
+++ b/unifi_client.py
@@ -158,8 +158,17 @@ class UnifiClient:
         logger.warning("Controller type detection failed, defaulting to UniFi OS")
 
     async def _authenticate(self) -> None:
-        """Authenticate with the UniFi controller."""
-        login_url = f"{self.base_url}{self.api_prefix}/api/login"
+        """Authenticate with the UniFi controller.
+
+        UniFi OS (UDM, Cloud Gateway) uses /api/auth/login
+        Standalone controllers use /api/login
+        """
+        if self._is_unifi_os:
+            # UniFi OS uses a different auth endpoint (no api_prefix needed)
+            login_url = f"{self.base_url}/api/auth/login"
+        else:
+            # Standalone controller
+            login_url = f"{self.base_url}/api/login"
 
         payload = {
             "username": self.username,