From d72504e580aba8c87b942cdbf3c2eafa92aab09e Mon Sep 17 00:00:00 2001
From: Punch <punch@punchthedev.ai>
Date: Thu, 4 Jun 2026 02:15:16 +0000
Subject: [PATCH] Add Swagger security schemes for admin and hidden routes

---
 app/main.py | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/app/main.py b/app/main.py
index 6b1b200..da8cd24 100644
--- a/app/main.py
+++ b/app/main.py
@@ -4,6 +4,7 @@
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.openapi.utils import get_openapi
 
 from app.db import init_db
 from app.routes import eval_preview, health, hidden, leaderboard, rounds, sota, specs, submissions
@@ -23,6 +24,45 @@ async def lifespan(app: FastAPI):
     lifespan=lifespan,
 )
 
+
+def custom_openapi():
+    if app.openapi_schema:
+        return app.openapi_schema
+    schema = get_openapi(
+        title=app.title,
+        version=app.version,
+        description=app.description,
+        routes=app.routes,
+    )
+    # Add security schemes so Swagger UI shows lock icons on protected endpoints
+    schema.setdefault("components", {})
+    schema["components"]["securitySchemes"] = {
+        "AdminToken": {
+            "type": "apiKey",
+            "in": "header",
+            "name": "X-Admin-Token",
+            "description": "Required for /admin/submissions/* endpoints",
+        },
+        "BearerToken": {
+            "type": "http",
+            "scheme": "bearer",
+            "description": "Required for /admin/hidden/* endpoints (FORGE_ADMIN_KEY)",
+        },
+    }
+    # Tag admin and hidden routes with their security requirements
+    for path, methods in schema.get("paths", {}).items():
+        for method, op in methods.items():
+            tags = op.get("tags", [])
+            if "admin" in tags:
+                op["security"] = [{"AdminToken": []}]
+            if "hidden" in tags:
+                op["security"] = [{"BearerToken": []}]
+    app.openapi_schema = schema
+    return schema
+
+
+app.openapi = custom_openapi
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],