feat: Add API key authentication and health endpoint

- Add API key authentication middleware with multiple auth methods (Bearer, X-API-Key, query param) - Add /health endpoint with server status, queue info, device info, and VRAM stats - Add CLI arguments --api-key and --api-key-file for authentication configuration - Static files and WebSocket connections exempt from authentication - Fully backward compatible - no authentication required by default - Add comprehensive documentation, examples, and test scripts
2025-12-18 02:23:06 +08:00 · 2025-12-11 15:33:08 +03:00 · 2025-12-11 15:33:08 +03:00 · 06bf79b19b
commit 06bf79b19b
parent fc657f471a
11 changed files with 1387 additions and 0 deletions
--- a/API_AUTHENTICATION.md
+++ b/API_AUTHENTICATION.md
@ -0,0 +1,221 @@
+# API Key Authentication and Health Check
+
+## Overview
+
+This implementation adds API key authentication protection to the ComfyUI REST API and a health check endpoint.
+
+## Features
+
+### 1. API Key Authentication
+
+Protects all API endpoints (except exempt ones) with API key authentication.
+
+#### Configuration
+
+You can enable API key authentication in two ways:
+
+**Option 1: Command line argument**
+```bash
+python main.py --api-key "your-secret-api-key-here"
+```
+
+**Option 2: API key file (more secure)**
+```bash
+# Create a file with your API key
+echo "your-secret-api-key-here" > api_key.txt
+
+# Start ComfyUI with the API key file
+python main.py --api-key-file api_key.txt
+```
+
+#### Using the API with Authentication
+
+When API key authentication is enabled, you must provide the API key in your requests:
+
+**Method 1: Authorization Header (Bearer Token)**
+```bash
+curl -H "Authorization: Bearer your-secret-api-key-here" http://localhost:8188/prompt
+```
+
+**Method 2: X-API-Key Header**
+```bash
+curl -H "X-API-Key: your-secret-api-key-here" http://localhost:8188/prompt
+```
+
+**Method 3: Query Parameter (less secure, for testing only)**
+```bash
+curl "http://localhost:8188/prompt?api_key=your-secret-api-key-here"
+```
+
+#### Exempt Endpoints
+
+The following endpoints do NOT require authentication:
+- `/health` - Health check endpoint
+- `/` - Root page (frontend)
+- `/ws` - WebSocket endpoint
+
+### 2. Health Check Endpoint
+
+A new `/health` endpoint provides server status information.
+
+#### Usage
+
+```bash
+curl http://localhost:8188/health
+```
+
+#### Response Format
+
+```json
+{
+  "status": "healthy",
+  "version": "0.4.0",
+  "timestamp": 1702307890.123,
+  "queue": {
+    "pending": 0,
+    "running": 0
+  },
+  "device": "cuda:0",
+  "vram": {
+    "total": 8589934592,
+    "free": 6442450944,
+    "used": 2147483648
+  }
+}
+```
+
+If the server is unhealthy, it returns a 503 status code:
+
+```json
+{
+  "status": "unhealthy",
+  "error": "error message here",
+  "timestamp": 1702307890.123
+}
+```
+
+## Examples
+
+### Starting ComfyUI with API Key Protection
+
+```bash
+# With direct API key
+python main.py --api-key "my-super-secret-key-12345"
+
+# With API key from file
+python main.py --api-key-file /path/to/api_key.txt
+
+# With API key and custom port
+python main.py --api-key "my-key" --port 8080
+```
+
+### Making Authenticated Requests
+
+**Python example:**
+```python
+import requests
+
+API_KEY = "your-api-key-here"
+BASE_URL = "http://localhost:8188"
+
+# Using Authorization header
+headers = {
+    "Authorization": f"Bearer {API_KEY}"
+}
+
+# Check health
+response = requests.get(f"{BASE_URL}/health")
+print(response.json())
+
+# Make authenticated request
+response = requests.post(
+    f"{BASE_URL}/prompt",
+    headers=headers,
+    json={"prompt": {...}}
+)
+print(response.json())
+```
+
+**JavaScript example:**
+```javascript
+const API_KEY = "your-api-key-here";
+const BASE_URL = "http://localhost:8188";
+
+// Using fetch with Authorization header
+async function makeRequest(endpoint, data) {
+  const response = await fetch(`${BASE_URL}${endpoint}`, {
+    method: 'POST',
+    headers: {
+      'Authorization': `Bearer ${API_KEY}`,
+      'Content-Type': 'application/json'
+    },
+    body: JSON.stringify(data)
+  });
+  return response.json();
+}
+
+// Check health (no auth required)
+fetch(`${BASE_URL}/health`)
+  .then(r => r.json())
+  .then(data => console.log(data));
+```
+
+### Monitoring with Health Check
+
+You can use the health endpoint for monitoring and health checks:
+
+```bash
+# Simple health check
+curl http://localhost:8188/health
+
+# Use in a monitoring script
+#!/bin/bash
+response=$(curl -s http://localhost:8188/health)
+status=$(echo $response | jq -r '.status')
+
+if [ "$status" == "healthy" ]; then
+  echo "✓ ComfyUI is healthy"
+  exit 0
+else
+  echo "✗ ComfyUI is unhealthy"
+  exit 1
+fi
+```
+
+## Security Considerations
+
+1. **Keep your API key secret**: Never commit API keys to version control
+2. **Use API key files**: Store API keys in separate files with restricted permissions
+3. **Use HTTPS in production**: Combine with `--tls-keyfile` and `--tls-certfile` options
+4. **Rotate keys regularly**: Change your API key periodically
+5. **Use strong keys**: Generate long, random API keys (e.g., using `openssl rand -hex 32`)
+
+### Generating a Secure API Key
+
+```bash
+# Generate a secure random API key
+openssl rand -hex 32
+
+# Or using Python
+python -c "import secrets; print(secrets.token_hex(32))"
+```
+
+## Troubleshooting
+
+### 401 Unauthorized Error
+
+If you receive a 401 error:
+- Verify the API key is correct
+- Check that you're including the key in the correct header format
+- Ensure there are no extra spaces or newlines in the key
+
+### Health Check Returns 503
+
+If the health check returns 503:
+- Check the server logs for error details
+- Verify ComfyUI started correctly
+- Check system resources (memory, disk space)
+
+## Disabling Authentication
+
+To disable API key authentication, simply don't provide the `--api-key` or `--api-key-file` arguments when starting ComfyUI. The server will work exactly as before with no authentication required.
--- a/API_SECURITY_IMPLEMENTATION.md
+++ b/API_SECURITY_IMPLEMENTATION.md
@ -0,0 +1,142 @@
+# ComfyUI API Security Enhancement
+
+## Summary
+
+This implementation adds API key authentication and a health check endpoint to ComfyUI.
+
+## Files Modified
+
+1. **middleware/auth_middleware.py** (NEW)
+   - API key authentication middleware
+   - Supports multiple authentication methods (Bearer token, X-API-Key header, query parameter)
+   - Configurable exempt paths
+
+2. **comfy/cli_args.py** (MODIFIED)
+   - Added `--api-key` argument for inline API key
+   - Added `--api-key-file` argument for API key from file
+   - Added logic to load API key from file
+
+3. **server.py** (MODIFIED)
+   - Imported auth middleware
+   - Integrated middleware into application
+   - Added `/health` endpoint with system information
+   - Configured exempt paths (/, /health, /ws)
+
+## New Files
+
+1. **API_AUTHENTICATION.md** - Complete documentation
+2. **test_api_auth.py** - Test suite for authentication
+3. **examples_api_auth.py** - Python usage examples
+
+## Quick Start
+
+### 1. Start ComfyUI with API Key Protection
+
+```bash
+# Generate a secure API key
+python -c "import secrets; print(secrets.token_hex(32))"
+
+# Start with API key
+python main.py --api-key "your-generated-key-here"
+
+# Or use a file
+echo "your-generated-key-here" > api_key.txt
+python main.py --api-key-file api_key.txt
+```
+
+### 2. Test the Health Endpoint
+
+```bash
+curl http://localhost:8188/health
+```
+
+### 3. Make Authenticated Requests
+
+```bash
+# Using Bearer token
+curl -H "Authorization: Bearer your-api-key" http://localhost:8188/prompt
+
+# Using X-API-Key header
+curl -H "X-API-Key: your-api-key" http://localhost:8188/prompt
+```
+
+### 4. Run Tests
+
+```bash
+# Install requests if needed
+pip install requests
+
+# Run test suite
+python test_api_auth.py your-api-key
+
+# Run examples
+python examples_api_auth.py
+```
+
+## Features
+
+### API Key Authentication
+- ✅ Multiple authentication methods (Bearer, X-API-Key, query param)
+- ✅ Configurable via command line
+- ✅ Secure file-based configuration
+- ✅ Exempt paths for health checks and WebSocket
+- ✅ Detailed logging of authentication attempts
+
+### Health Check Endpoint
+- ✅ Returns server status
+- ✅ Queue information (pending/running)
+- ✅ Device information
+- ✅ VRAM usage (if GPU available)
+- ✅ Version information
+- ✅ Timestamp for monitoring
+
+## Security Best Practices
+
+1. **Generate Strong Keys**: Use `openssl rand -hex 32` or similar
+2. **Use File-Based Config**: Keep keys out of command history
+3. **Enable HTTPS**: Use with `--tls-keyfile` and `--tls-certfile`
+4. **Restrict File Permissions**: `chmod 600 api_key.txt`
+5. **Rotate Keys Regularly**: Change API keys periodically
+6. **Monitor Access**: Check logs for unauthorized attempts
+
+## Backward Compatibility
+
+- ✅ Fully backward compatible
+- ✅ No authentication required by default
+- ✅ Existing functionality unchanged
+- ✅ WebSocket connections work normally
+
+## Testing
+
+The implementation has been tested for:
+- ✅ Syntax errors (none found)
+- ✅ Import compatibility
+- ✅ Middleware integration
+- ✅ Route configuration
+- ✅ Health endpoint functionality
+
+To fully test in your environment:
+```bash
+# 1. Start server without auth (test backward compatibility)
+python main.py
+
+# 2. Start server with auth
+python main.py --api-key "test-key-123"
+
+# 3. Run test suite
+python test_api_auth.py test-key-123
+
+# 4. Check health endpoint
+curl http://localhost:8188/health
+```
+
+## Support
+
+For detailed documentation, see:
+- **API_AUTHENTICATION.md** - Complete usage guide
+- **examples_api_auth.py** - Code examples
+- **test_api_auth.py** - Test suite
+
+## License
+
+Same as ComfyUI main project.
--- a/QUICK_START_AUTH.md
+++ b/QUICK_START_AUTH.md
@ -0,0 +1,118 @@
+# Quick Start Guide - API Authentication
+
+## Step-by-Step Instructions
+
+### 1. Start ComfyUI with API Key
+
+```bash
+# Stop any running ComfyUI instance first
+# Then start with an API key:
+
+python main.py --api-key "my-secret-key-123"
+```
+
+**You should see in the logs:**
+```
+[Auth] API Key authentication enabled
+```
+
+### 2. Test the Authentication
+
+**Health check (works without auth):**
+```bash
+curl http://localhost:8188/health
+```
+
+**Protected endpoint without auth (should fail):**
+```bash
+curl http://localhost:8188/object_info
+# Should return: {"error": "Unauthorized", "message": "..."}
+```
+
+**Protected endpoint with auth (should work):**
+```bash
+curl -H "Authorization: Bearer my-secret-key-123" http://localhost:8188/object_info
+# Should return: {...node definitions...}
+```
+
+### 3. Run the Test Script
+
+```bash
+chmod +x test_auth_quick.sh
+./test_auth_quick.sh
+```
+
+## Common Issues
+
+### Issue: All requests work without authentication
+
+**Problem:** You didn't start the server with `--api-key`
+
+**Solution:**
+```bash
+# Stop the server (Ctrl+C)
+# Restart with API key:
+python main.py --api-key "your-key-here"
+```
+
+**Verify it's enabled:**
+```bash
+# In another terminal, check if auth is working:
+curl http://localhost:8188/object_info
+# Should return 401 Unauthorized
+```
+
+### Issue: Authentication is enabled but I get 401 even with correct key
+
+**Problem:** Key format or typo
+
+**Solution:**
+- Ensure no extra spaces in the key
+- Check the Authorization header format: `Authorization: Bearer YOUR_KEY`
+- Try X-API-Key header: `X-API-Key: YOUR_KEY`
+
+## Example: Full Workflow
+
+```bash
+# 1. Generate a secure key
+python -c "import secrets; print(secrets.token_hex(32))"
+# Output: a1b2c3d4e5f6...
+
+# 2. Save to file
+echo "a1b2c3d4e5f6..." > api_key.txt
+
+# 3. Start server with key file
+python main.py --api-key-file api_key.txt
+
+# 4. Use the API
+API_KEY=$(cat api_key.txt)
+curl -H "Authorization: Bearer $API_KEY" http://localhost:8188/object_info
+```
+
+## Test with Python
+
+```python
+import requests
+
+API_KEY = "my-secret-key-123"
+BASE_URL = "http://localhost:8188"
+
+# This should fail (no auth)
+response = requests.get(f"{BASE_URL}/object_info")
+print(f"No auth: {response.status_code}")  # Should be 401
+
+# This should work (with auth)
+headers = {"Authorization": f"Bearer {API_KEY}"}
+response = requests.get(f"{BASE_URL}/object_info", headers=headers)
+print(f"With auth: {response.status_code}")  # Should be 200
+```
+
+## Disable Authentication
+
+Simply start ComfyUI without the `--api-key` argument:
+
+```bash
+python main.py
+```
+
+The server will work exactly as before with no authentication required.
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@ -42,6 +42,9 @@ parser.add_argument("--tls-certfile", type=str, help="Path to TLS (SSL) certific
 parser.add_argument("--enable-cors-header", type=str, default=None, metavar="ORIGIN", nargs="?", const="*", help="Enable CORS (Cross-Origin Resource Sharing) with optional origin or allow all with default '*'.")
 parser.add_argument("--max-upload-size", type=float, default=100, help="Set the maximum upload size in MB.")

+parser.add_argument("--api-key", type=str, default=None, help="Require API key authentication for all API endpoints except health check. Provide the key via 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header.")
+parser.add_argument("--api-key-file", type=str, default=None, help="Path to a file containing the API key. Alternative to --api-key for better security.")
+
 parser.add_argument("--base-directory", type=str, default=None, help="Set the ComfyUI base directory for models, custom_nodes, input, output, temp, and user directories.")
 parser.add_argument("--extra-model-paths-config", type=str, default=None, metavar="PATH", nargs='+', action='append', help="Load one or more extra_model_paths.yaml files.")
 parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory. Overrides --base-directory.")
@ -239,6 +242,14 @@ if args.disable_auto_launch:
 if args.force_fp16:
    args.fp16_unet = True

+# Load API key from file if specified
+if args.api_key_file and not args.api_key:
+    try:
+        with open(args.api_key_file, 'r') as f:
+            args.api_key = f.read().strip()
+    except Exception as e:
+        print(f"Error reading API key from file {args.api_key_file}: {e}")
+        args.api_key = None

 # '--fast' is not provided, use an empty set
 if args.fast is None:
--- a/examples_api_auth.py
+++ b/examples_api_auth.py
@ -0,0 +1,204 @@
+"""
+Example: Using ComfyUI API with Authentication
+"""
+
+import requests
+import json
+
+# Your API configuration
+API_KEY = "your-api-key-here"
+BASE_URL = "http://localhost:8188"
+
+
+def example_health_check():
+    """Example: Check server health (no authentication required)"""
+    print("=== Health Check Example ===")
+    
+    response = requests.get(f"{BASE_URL}/health")
+    
+    if response.status_code == 200:
+        health = response.json()
+        print(f"Status: {health['status']}")
+        print(f"Version: {health['version']}")
+        print(f"Queue - Pending: {health['queue']['pending']}, Running: {health['queue']['running']}")
+        if 'device' in health:
+            print(f"Device: {health['device']}")
+        if 'vram' in health:
+            vram = health['vram']
+            vram_used_gb = vram['used'] / (1024**3)
+            vram_total_gb = vram['total'] / (1024**3)
+            print(f"VRAM: {vram_used_gb:.2f} GB / {vram_total_gb:.2f} GB")
+    else:
+        print(f"Health check failed with status {response.status_code}")
+    
+    print()
+
+
+def example_get_object_info():
+    """Example: Get object info with authentication"""
+    print("=== Get Object Info Example ===")
+    
+    # Method 1: Using Authorization Bearer header
+    headers = {
+        "Authorization": f"Bearer {API_KEY}"
+    }
+    
+    response = requests.get(f"{BASE_URL}/object_info", headers=headers)
+    
+    if response.status_code == 200:
+        print("✓ Successfully retrieved object info")
+        object_info = response.json()
+        print(f"Number of node types: {len(object_info)}")
+    elif response.status_code == 401:
+        print("✗ Authentication failed - check your API key")
+        print(response.json())
+    else:
+        print(f"✗ Request failed with status {response.status_code}")
+    
+    print()
+
+
+def example_queue_prompt():
+    """Example: Queue a prompt with authentication"""
+    print("=== Queue Prompt Example ===")
+    
+    # Simple workflow example
+    workflow = {
+        "prompt": {
+            "1": {
+                "inputs": {
+                    "text": "a beautiful landscape"
+                },
+                "class_type": "CLIPTextEncode"
+            }
+        },
+        "client_id": "example_client"
+    }
+    
+    # Using Authorization Bearer header
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
+    }
+    
+    response = requests.post(
+        f"{BASE_URL}/prompt",
+        headers=headers,
+        json=workflow
+    )
+    
+    if response.status_code == 200:
+        result = response.json()
+        print("✓ Prompt queued successfully")
+        print(f"Prompt ID: {result.get('prompt_id', 'N/A')}")
+    elif response.status_code == 401:
+        print("✗ Authentication failed - check your API key")
+        print(response.json())
+    else:
+        print(f"✗ Request failed with status {response.status_code}")
+        print(response.text)
+    
+    print()
+
+
+def example_using_session():
+    """Example: Using requests.Session for multiple requests"""
+    print("=== Session Example (Multiple Requests) ===")
+    
+    # Create a session with authentication header
+    session = requests.Session()
+    session.headers.update({
+        "Authorization": f"Bearer {API_KEY}"
+    })
+    
+    # Now all requests will automatically include the auth header
+    
+    # Request 1: Get embeddings
+    response = session.get(f"{BASE_URL}/embeddings")
+    if response.status_code == 200:
+        print(f"✓ Got embeddings list")
+    
+    # Request 2: Get queue
+    response = session.get(f"{BASE_URL}/queue")
+    if response.status_code == 200:
+        queue = response.json()
+        print(f"✓ Got queue info - Pending: {len(queue.get('queue_pending', []))}")
+    
+    # Request 3: Get system stats
+    response = session.get(f"{BASE_URL}/system_stats")
+    if response.status_code == 200:
+        print(f"✓ Got system stats")
+    
+    print()
+
+
+def example_error_handling():
+    """Example: Proper error handling"""
+    print("=== Error Handling Example ===")
+    
+    headers = {
+        "Authorization": f"Bearer {API_KEY}"
+    }
+    
+    try:
+        response = requests.get(f"{BASE_URL}/queue", headers=headers, timeout=5)
+        response.raise_for_status()  # Raises exception for 4xx/5xx status codes
+        
+        data = response.json()
+        print("✓ Request successful")
+        print(f"Queue pending: {len(data.get('queue_pending', []))}")
+        print(f"Queue running: {len(data.get('queue_running', []))}")
+        
+    except requests.exceptions.Timeout:
+        print("✗ Request timed out")
+    except requests.exceptions.ConnectionError:
+        print("✗ Could not connect to server")
+    except requests.exceptions.HTTPError as e:
+        if e.response.status_code == 401:
+            print("✗ Authentication failed - invalid API key")
+        elif e.response.status_code == 403:
+            print("✗ Access forbidden")
+        else:
+            print(f"✗ HTTP error: {e}")
+    except Exception as e:
+        print(f"✗ Unexpected error: {e}")
+    
+    print()
+
+
+# Alternative authentication methods
+def example_alternative_auth_methods():
+    """Example: Different ways to provide API key"""
+    print("=== Alternative Authentication Methods ===")
+    
+    # Method 1: Authorization Bearer token (recommended)
+    headers1 = {"Authorization": f"Bearer {API_KEY}"}
+    response1 = requests.get(f"{BASE_URL}/embeddings", headers=headers1)
+    print(f"Method 1 (Bearer): Status {response1.status_code}")
+    
+    # Method 2: X-API-Key header
+    headers2 = {"X-API-Key": API_KEY}
+    response2 = requests.get(f"{BASE_URL}/embeddings", headers=headers2)
+    print(f"Method 2 (X-API-Key): Status {response2.status_code}")
+    
+    # Method 3: Query parameter (less secure, not recommended for production)
+    response3 = requests.get(f"{BASE_URL}/embeddings?api_key={API_KEY}")
+    print(f"Method 3 (Query param): Status {response3.status_code}")
+    
+    print()
+
+
+if __name__ == "__main__":
+    print("ComfyUI API Authentication Examples")
+    print("=" * 60)
+    print()
+    
+    # Run examples
+    example_health_check()
+    example_get_object_info()
+    example_using_session()
+    example_error_handling()
+    example_alternative_auth_methods()
+    
+    print("=" * 60)
+    print("All examples completed!")
--- a/generate_api_key.sh
+++ b/generate_api_key.sh
@ -0,0 +1,75 @@
+#!/bin/bash
+
+# ComfyUI API Key Generator
+# This script helps you generate and configure API keys for ComfyUI
+
+set -e
+
+echo "================================================"
+echo "ComfyUI API Key Generator"
+echo "================================================"
+echo ""
+
+# Function to generate a random API key
+generate_key() {
+    if command -v openssl >/dev/null 2>&1; then
+        openssl rand -hex 32
+    elif command -v python3 >/dev/null 2>&1; then
+        python3 -c "import secrets; print(secrets.token_hex(32))"
+    elif command -v python >/dev/null 2>&1; then
+        python -c "import secrets; print(secrets.token_hex(32))"
+    else
+        echo "Error: Neither openssl nor python is available to generate random key"
+        exit 1
+    fi
+}
+
+# Generate the API key
+echo "Generating secure API key..."
+API_KEY=$(generate_key)
+echo ""
+echo "Generated API Key:"
+echo "================================================"
+echo "$API_KEY"
+echo "================================================"
+echo ""
+
+# Ask user if they want to save to file
+read -p "Would you like to save this key to a file? (y/n) " -n 1 -r
+echo ""
+
+if [[ $REPLY =~ ^[Yy]$ ]]; then
+    # Get filename
+    read -p "Enter filename (default: api_key.txt): " FILENAME
+    FILENAME=${FILENAME:-api_key.txt}
+    
+    # Save the key
+    echo "$API_KEY" > "$FILENAME"
+    
+    # Set restrictive permissions
+    chmod 600 "$FILENAME"
+    
+    echo "✓ API key saved to: $FILENAME"
+    echo "✓ File permissions set to 600 (owner read/write only)"
+    echo ""
+    echo "To start ComfyUI with this API key:"
+    echo "  python main.py --api-key-file $FILENAME"
+else
+    echo ""
+    echo "To start ComfyUI with this API key:"
+    echo "  python main.py --api-key \"$API_KEY\""
+fi
+
+echo ""
+echo "================================================"
+echo "Important Security Notes:"
+echo "================================================"
+echo "1. Keep this key secret - don't commit it to git"
+echo "2. Use HTTPS in production for encrypted transport"
+echo "3. Rotate keys regularly"
+echo "4. Add your key file to .gitignore"
+echo ""
+echo "Example .gitignore entry:"
+echo "  api_key.txt"
+echo "  *.key"
+echo "================================================"
--- a/middleware/auth_middleware.py
+++ b/middleware/auth_middleware.py
@ -0,0 +1,139 @@
+"""API Key Authentication middleware for ComfyUI server"""
+
+from aiohttp import web
+from typing import Callable, Awaitable, Optional, Set
+import logging
+import os
+
+
+class APIKeyAuth:
+    """API Key Authentication handler"""
+    
+    def __init__(self, api_key: Optional[str] = None, exempt_paths: Optional[Set[str]] = None):
+        """
+        Initialize API Key Authentication
+        
+        Args:
+            api_key: The API key to validate against. If None, authentication is disabled.
+            exempt_paths: Set of paths that don't require authentication (e.g., health check)
+        """
+        self.api_key = api_key
+        self.enabled = api_key is not None and len(api_key) > 0
+        self.exempt_paths = exempt_paths or {"/health"}
+        
+        # Static file extensions that don't require authentication
+        self.static_extensions = {
+            '.html', '.js', '.css', '.json', '.map', '.png', '.jpg', '.jpeg', 
+            '.gif', '.svg', '.ico', '.woff', '.woff2', '.ttf', '.eot', '.webp'
+        }
+        
+        # Path prefixes that serve static content
+        self.static_path_prefixes = {
+            '/extensions/', '/templates/', '/docs/'
+        }
+        
+        if self.enabled:
+            logging.info("[Auth] API Key authentication enabled")
+        else:
+            logging.info("[Auth] API Key authentication disabled")
+    
+    def is_path_exempt(self, path: str) -> bool:
+        """Check if a path is exempt from authentication"""
+        # Exact match for specific exempt paths
+        if path in self.exempt_paths:
+            return True
+        
+        # Root path for index.html
+        if path == "/":
+            return True
+        
+        # Static file extensions
+        for ext in self.static_extensions:
+            if path.endswith(ext):
+                return True
+        
+        # Static path prefixes (extensions, templates, docs, etc.)
+        for prefix in self.static_path_prefixes:
+            if path.startswith(prefix):
+                return True
+        
+        return False
+    
+    def validate_api_key(self, provided_key: Optional[str]) -> bool:
+        """Validate the provided API key"""
+        if not self.enabled:
+            return True
+        
+        if not provided_key:
+            return False
+        
+        return provided_key == self.api_key
+    
+    def extract_api_key(self, request: web.Request) -> Optional[str]:
+        """
+        Extract API key from request.
+        Checks Authorization header (Bearer token) and X-API-Key header.
+        """
+        # Check Authorization header (Bearer token)
+        auth_header = request.headers.get("Authorization", "")
+        if auth_header.startswith("Bearer "):
+            return auth_header[7:]  # Remove "Bearer " prefix
+        
+        # Check X-API-Key header
+        api_key_header = request.headers.get("X-API-Key", "")
+        if api_key_header:
+            return api_key_header
+        
+        # Check query parameter (less secure, but convenient for testing)
+        api_key_query = request.query.get("api_key", "")
+        if api_key_query:
+            return api_key_query
+        
+        return None
+
+
+def create_api_key_middleware(api_key: Optional[str] = None, exempt_paths: Optional[Set[str]] = None):
+    """
+    Create API key authentication middleware
+    
+    Args:
+        api_key: The API key to validate against. If None, authentication is disabled.
+        exempt_paths: Set of paths that don't require authentication
+    
+    Returns:
+        Middleware function for aiohttp
+    """
+    auth = APIKeyAuth(api_key, exempt_paths)
+    
+    @web.middleware
+    async def api_key_middleware(
+        request: web.Request, 
+        handler: Callable[[web.Request], Awaitable[web.Response]]
+    ) -> web.Response:
+        """Middleware to validate API key for protected endpoints"""
+        
+        # Skip authentication if disabled
+        if not auth.enabled:
+            return await handler(request)
+        
+        # Check if path is exempt from authentication
+        if auth.is_path_exempt(request.path):
+            return await handler(request)
+        
+        # Extract and validate API key
+        provided_key = auth.extract_api_key(request)
+        
+        if not auth.validate_api_key(provided_key):
+            logging.warning(f"[Auth] Unauthorized access attempt to {request.path} from {request.remote}")
+            return web.json_response(
+                {
+                    "error": "Unauthorized",
+                    "message": "Invalid or missing API key. Provide API key via 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header."
+                },
+                status=401
+            )
+        
+        # API key is valid, proceed with request
+        return await handler(request)
+    
+    return api_key_middleware
--- a/server.py
+++ b/server.py
@ -43,6 +43,7 @@ from protocol import BinaryEventTypes

 # Import cache control middleware
 from middleware.cache_middleware import cache_control
+from middleware.auth_middleware import create_api_key_middleware

 if args.enable_manager:
    import comfyui_manager
@ -204,6 +205,17 @@ class PromptServer():
        self.number = 0

        middlewares = [cache_control, deprecation_warning]
+        
+        # Add API key authentication middleware if enabled
+        if args.api_key:
+            # Define paths that don't require authentication
+            # Note: Static files (.js, .css, .html, etc.) and root "/" are automatically exempted
+            exempt_paths = {
+                "/health",  # Health check endpoint
+                "/ws",  # WebSocket endpoint
+            }
+            middlewares.append(create_api_key_middleware(args.api_key, exempt_paths))
+        
        if args.enable_compress_response_body:
            middlewares.append(compress_body)

@ -303,6 +315,50 @@ class PromptServer():
            response.headers["Expires"] = "0"
            return response

+        @routes.get("/health")
+        async def get_health(request):
+            """Health check endpoint that returns the status of the server"""
+            try:
+                # Basic health information
+                health_data = {
+                    "status": "healthy",
+                    "version": __version__,
+                    "timestamp": time.time(),
+                    "queue": {
+                        "pending": len(self.prompt_queue.queue),
+                        "running": len(self.prompt_queue.currently_running)
+                    }
+                }
+                
+                # Add device info if available
+                try:
+                    device = comfy.model_management.get_torch_device()
+                    health_data["device"] = str(device)
+                    
+                    # Add VRAM info if GPU is available
+                    if comfy.model_management.vram_state != comfy.model_management.VRAMState.DISABLED:
+                        vram_total = comfy.model_management.get_total_memory()
+                        vram_free = comfy.model_management.get_free_memory()
+                        health_data["vram"] = {
+                            "total": vram_total,
+                            "free": vram_free,
+                            "used": vram_total - vram_free
+                        }
+                except Exception as e:
+                    logging.debug(f"Could not get device info for health check: {e}")
+                
+                return web.json_response(health_data)
+            except Exception as e:
+                logging.error(f"Health check failed: {e}")
+                return web.json_response(
+                    {
+                        "status": "unhealthy",
+                        "error": str(e),
+                        "timestamp": time.time()
+                    },
+                    status=503
+                )
+
        @routes.get("/embeddings")
        def get_embeddings(request):
            embeddings = folder_paths.get_filename_list("embeddings")
--- a/test_api_auth.py
+++ b/test_api_auth.py
@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+"""
+Test script for ComfyUI API Key Authentication and Health Check
+
+This script demonstrates how to:
+1. Check the health endpoint (no auth required)
+2. Make authenticated requests to the API
+"""
+
+import requests
+import json
+import sys
+
+# Configuration
+BASE_URL = "http://localhost:8188"
+API_KEY = "your-api-key-here"  # Replace with your actual API key
+
+
+def test_health_check():
+    """Test the health check endpoint (no authentication required)"""
+    print("Testing health check endpoint...")
+    try:
+        response = requests.get(f"{BASE_URL}/health")
+        print(f"Status Code: {response.status_code}")
+        print(f"Response: {json.dumps(response.json(), indent=2)}")
+        return response.status_code == 200
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def test_without_auth():
+    """Test accessing protected endpoint without authentication"""
+    print("\nTesting access without authentication...")
+    try:
+        response = requests.get(f"{BASE_URL}/object_info")
+        print(f"Status Code: {response.status_code}")
+        if response.status_code == 401:
+            print("✓ Correctly rejected (401 Unauthorized)")
+            print(f"Response: {json.dumps(response.json(), indent=2)}")
+            return True
+        elif response.status_code == 200:
+            print("✓ No authentication required (API key not enabled)")
+            return True
+        else:
+            print(f"✗ Unexpected status code: {response.status_code}")
+            return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def test_with_bearer_token():
+    """Test accessing protected endpoint with Bearer token"""
+    print("\nTesting with Bearer token authentication...")
+    try:
+        headers = {
+            "Authorization": f"Bearer {API_KEY}"
+        }
+        response = requests.get(f"{BASE_URL}/object_info", headers=headers)
+        print(f"Status Code: {response.status_code}")
+        if response.status_code == 200:
+            print("✓ Successfully authenticated with Bearer token")
+            return True
+        elif response.status_code == 401:
+            print("✗ Authentication failed (check your API key)")
+            print(f"Response: {json.dumps(response.json(), indent=2)}")
+            return False
+        else:
+            print(f"✗ Unexpected status code: {response.status_code}")
+            return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def test_with_api_key_header():
+    """Test accessing protected endpoint with X-API-Key header"""
+    print("\nTesting with X-API-Key header authentication...")
+    try:
+        headers = {
+            "X-API-Key": API_KEY
+        }
+        response = requests.get(f"{BASE_URL}/object_info", headers=headers)
+        print(f"Status Code: {response.status_code}")
+        if response.status_code == 200:
+            print("✓ Successfully authenticated with X-API-Key header")
+            return True
+        elif response.status_code == 401:
+            print("✗ Authentication failed (check your API key)")
+            print(f"Response: {json.dumps(response.json(), indent=2)}")
+            return False
+        else:
+            print(f"✗ Unexpected status code: {response.status_code}")
+            return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def test_with_query_parameter():
+    """Test accessing protected endpoint with query parameter"""
+    print("\nTesting with query parameter authentication...")
+    try:
+        response = requests.get(f"{BASE_URL}/object_info?api_key={API_KEY}")
+        print(f"Status Code: {response.status_code}")
+        if response.status_code == 200:
+            print("✓ Successfully authenticated with query parameter")
+            return True
+        elif response.status_code == 401:
+            print("✗ Authentication failed (check your API key)")
+            print(f"Response: {json.dumps(response.json(), indent=2)}")
+            return False
+        else:
+            print(f"✗ Unexpected status code: {response.status_code}")
+            return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def main():
+    """Run all tests"""
+    print("=" * 60)
+    print("ComfyUI API Authentication Test Suite")
+    print("=" * 60)
+    print(f"Base URL: {BASE_URL}")
+    print(f"API Key: {'*' * (len(API_KEY) - 4) + API_KEY[-4:] if len(API_KEY) > 4 else '***'}")
+    print("=" * 60)
+    
+    results = []
+    
+    # Test 1: Health check (always works)
+    results.append(("Health Check", test_health_check()))
+    
+    # Test 2: Without authentication (should fail if auth is enabled)
+    results.append(("No Auth", test_without_auth()))
+    
+    # Test 3: Bearer token authentication
+    results.append(("Bearer Token", test_with_bearer_token()))
+    
+    # Test 4: X-API-Key header authentication
+    results.append(("X-API-Key Header", test_with_api_key_header()))
+    
+    # Test 5: Query parameter authentication
+    results.append(("Query Parameter", test_with_query_parameter()))
+    
+    # Summary
+    print("\n" + "=" * 60)
+    print("Test Summary")
+    print("=" * 60)
+    for test_name, passed in results:
+        status = "✓ PASS" if passed else "✗ FAIL"
+        print(f"{test_name:20s} {status}")
+    
+    total = len(results)
+    passed = sum(1 for _, result in results if result)
+    print("=" * 60)
+    print(f"Total: {passed}/{total} tests passed")
+    print("=" * 60)
+    
+    # Exit with appropriate code
+    sys.exit(0 if passed == total else 1)
+
+
+if __name__ == "__main__":
+    # Check if user wants to override the API key
+    if len(sys.argv) > 1:
+        API_KEY = sys.argv[1]
+    
+    if API_KEY == "your-api-key-here":
+        print("WARNING: Using default API key. Set your API key as the first argument:")
+        print(f"  python {sys.argv[0]} YOUR_API_KEY")
+        print("")
+    
+    main()
--- a/test_auth_quick.sh
+++ b/test_auth_quick.sh
@ -0,0 +1,128 @@
+#!/bin/bash
+
+# Quick Test Script for ComfyUI API Authentication
+# This script tests that authentication is working correctly
+
+set -e
+
+API_KEY="test-key-123"
+BASE_URL="http://localhost:8188"
+
+echo "================================================"
+echo "ComfyUI API Authentication Test"
+echo "================================================"
+echo ""
+echo "IMPORTANT: Make sure ComfyUI is running with:"
+echo "  python main.py --api-key \"$API_KEY\""
+echo ""
+echo "Press Enter to continue or Ctrl+C to cancel..."
+read
+
+echo ""
+echo "================================================"
+echo "Test 1: Health endpoint (should work without auth)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" "$BASE_URL/health")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "200" ]; then
+    echo "✓ PASS - Health endpoint accessible without auth"
+else
+    echo "✗ FAIL - Health endpoint should return 200"
+fi
+echo ""
+
+echo "================================================"
+echo "Test 2: Protected endpoint without auth (should fail)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" "$BASE_URL/object_info")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "401" ]; then
+    echo "✓ PASS - Correctly rejected without auth"
+    echo "Response: $body"
+else
+    echo "✗ FAIL - Should return 401 Unauthorized"
+    echo "Response: $body"
+fi
+echo ""
+
+echo "================================================"
+echo "Test 3: Protected endpoint with wrong key (should fail)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+    -H "Authorization: Bearer wrong-key-456" \
+    "$BASE_URL/object_info")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "401" ]; then
+    echo "✓ PASS - Correctly rejected wrong key"
+    echo "Response: $body"
+else
+    echo "✗ FAIL - Should return 401 Unauthorized"
+    echo "Response: $body"
+fi
+echo ""
+
+echo "================================================"
+echo "Test 4: Protected endpoint with correct key (should work)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+    -H "Authorization: Bearer $API_KEY" \
+    "$BASE_URL/object_info")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "200" ]; then
+    echo "✓ PASS - Successfully authenticated"
+else
+    echo "✗ FAIL - Should return 200 OK"
+    echo "Response: $body"
+fi
+echo ""
+
+echo "================================================"
+echo "Test 5: X-API-Key header method (should work)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+    -H "X-API-Key: $API_KEY" \
+    "$BASE_URL/embeddings")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "200" ]; then
+    echo "✓ PASS - X-API-Key header works"
+else
+    echo "✗ FAIL - Should return 200 OK"
+    echo "Response: $body"
+fi
+echo ""
+
+echo "================================================"
+echo "Test 6: Query parameter method (should work)"
+echo "================================================"
+response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+    "$BASE_URL/embeddings?api_key=$API_KEY")
+status=$(echo "$response" | grep HTTP_STATUS | cut -d: -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "Status: $status"
+if [ "$status" = "200" ]; then
+    echo "✓ PASS - Query parameter works"
+else
+    echo "✗ FAIL - Should return 200 OK"
+    echo "Response: $body"
+fi
+echo ""
+
+echo "================================================"
+echo "All tests completed!"
+echo "================================================"
--- a/test_vibevoice_workflow.sh
+++ b/test_vibevoice_workflow.sh
@ -0,0 +1,117 @@
+#!/bin/bash
+
+# Test ComfyUI API with VibeVoice workflow
+# Usage: ./test_vibevoice_workflow.sh [API_KEY]
+
+# Configuration
+BASE_URL="http://localhost:8188"
+API_KEY="${1:-}"
+
+# Set headers based on whether API key is provided
+if [ -n "$API_KEY" ]; then
+    AUTH_HEADER="Authorization: Bearer $API_KEY"
+    echo "Using API Key authentication"
+else
+    AUTH_HEADER=""
+    echo "No API Key provided (running without authentication)"
+fi
+
+# The workflow payload
+# This converts the ComfyUI workflow format to the prompt API format
+read -r -d '' PAYLOAD << 'EOF'
+{
+  "prompt": {
+    "1": {
+      "inputs": {
+        "speaker_1_voice": ["2", 0],
+        "speaker_2_voice": null,
+        "speaker_3_voice": null,
+        "speaker_4_voice": null,
+        "model_name": "VibeVoice-Large",
+        "text": "[1] And this is a generated voice, how cool is that?",
+        "quantize_llm_4bit": false,
+        "attention_mode": "sdpa",
+        "cfg_scale": 1.3,
+        "inference_steps": 10,
+        "seed": 1117544514407045,
+        "do_sample": true,
+        "temperature": 0.95,
+        "top_p": 0.95,
+        "top_k": 0,
+        "force_offload": false
+      },
+      "class_type": "VibeVoiceTTS"
+    },
+    "2": {
+      "inputs": {
+        "audio": "audio1.wav"
+      },
+      "class_type": "LoadAudio"
+    },
+    "3": {
+      "inputs": {
+        "audio": ["1", 0],
+        "filename_prefix": "audio/ComfyUI"
+      },
+      "class_type": "SaveAudio"
+    }
+  },
+  "client_id": "test_client_$(date +%s)"
+}
+EOF
+
+echo ""
+echo "================================================"
+echo "Sending workflow to ComfyUI..."
+echo "================================================"
+echo ""
+
+# Make the request
+if [ -n "$AUTH_HEADER" ]; then
+    response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+        -X POST \
+        -H "Content-Type: application/json" \
+        -H "$AUTH_HEADER" \
+        -d "$PAYLOAD" \
+        "$BASE_URL/prompt")
+else
+    response=$(curl -s -w "\nHTTP_STATUS:%{http_code}" \
+        -X POST \
+        -H "Content-Type: application/json" \
+        -d "$PAYLOAD" \
+        "$BASE_URL/prompt")
+fi
+
+# Extract HTTP status
+http_status=$(echo "$response" | grep "HTTP_STATUS" | cut -d':' -f2)
+body=$(echo "$response" | sed '/HTTP_STATUS/d')
+
+echo "HTTP Status: $http_status"
+echo ""
+echo "Response:"
+echo "$body" | python3 -m json.tool 2>/dev/null || echo "$body"
+echo ""
+
+if [ "$http_status" = "200" ]; then
+    echo "✓ Workflow queued successfully!"
+    
+    # Extract prompt_id if available
+    prompt_id=$(echo "$body" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('prompt_id', ''))" 2>/dev/null)
+    if [ -n "$prompt_id" ]; then
+        echo "Prompt ID: $prompt_id"
+        echo ""
+        echo "To check status:"
+        if [ -n "$AUTH_HEADER" ]; then
+            echo "  curl -H \"$AUTH_HEADER\" $BASE_URL/history/$prompt_id"
+        else
+            echo "  curl $BASE_URL/history/$prompt_id"
+        fi
+    fi
+elif [ "$http_status" = "401" ]; then
+    echo "✗ Authentication failed - check your API key"
+else
+    echo "✗ Request failed"
+fi
+
+echo ""
+echo "================================================"