ComfyUI MCP Server

# Refactoring Implementation Summary ## ✅ Completed Work ### 1. Authentication Module (`src/auth/`) Created a clean authentication abstraction following SOLID principles: - **[src/auth/base.py](src/auth/base.py)**: Abstract `ComfyAuth` base class with `auth_header` property - **Implementations**: - `NoAuth`: Current default (no authentication) - `TokenAuth`: Future-ready Bearer token authentication - **Benefits**: - Dependency Inversion Principle: Routes depend on abstraction, not concrete auth - Easy to extend with new auth methods (API keys, OAuth, etc.) - Centralized header injection ### 2. HTTP Client Module (`src/client/`) Replaced synchronous `requests` library with async `httpx`: - **[src/client/get_data.py](src/client/get_data.py)**: Core async HTTP function - Authentication header injection via `ComfyAuth` - Session context manager for connection pooling - Standardized error handling - Support for GET, POST, HEAD, DELETE methods - **[src/client/response.py](src/client/response.py)**: Standardized response object - `ResponseGetData` dataclass for consistent returns - Status code, parsed response, success flag, headers - **Benefits**: - Async support for better concurrency - Connection pooling reduces latency - Consistent interface across all HTTP operations - Testable with httpx mocking ### 3. Logging Utilities (`src/utils/`) Decorator-based logging with correlation ID tracking and parameter sanitization: - **[src/utils/logging.py](src/utils/logging.py)**: `@log_call` decorator - Automatic entry/exit/timing logging with correlation IDs - Optional parameter logging with sensitive data redaction (`log_params`, `sensitive_params`) - Supports both sync and async functions - Zero-boilerplate correlation via `get_global_logger()` - No manual logging clutter in route functions - **[src/utils/context.py](src/utils/context.py)**: Correlation ID context management - `contextvars`-based request tracking - Auto-propagation across async boundaries - Used by `CorrelationLogger` for automatic ID injection - **Benefits**: - Aspect-Oriented Programming (AOP) pattern - End-to-end request tracing across all layers - Security: Automatic redaction of sensitive parameters (API keys, tokens, passwords) - Easy debugging with timing and correlation information - Reduced boilerplate in business logic ### 4. Route Functions (`src/routes/`) Thin HTTP transport layer with minimal logic: - **[src/routes/workflow.py](src/routes/workflow.py)**: - `queue_workflow()`: Submit workflow for execution - `get_history()`: Get all execution history - `get_prompt_history()`: Get specific prompt history - **[src/routes/queue.py](src/routes/queue.py)**: - `get_queue()`: Get current queue status - `cancel_prompt()`: Cancel queued/running prompt - **[src/routes/assets.py](src/routes/assets.py)**: - `get_asset()`: Download generated assets - `get_asset_metadata()`: Get asset headers (HEAD request) - **[src/routes/models.py](src/routes/models.py)**: - `get_available_models()`: Discover available checkpoints/models - **Benefits**: - Single Responsibility: Only HTTP transport - All routes follow identical pattern - Easy to test with httpx mocking - Custom exceptions per module - Comprehensive docstrings with examples ### 5. Documentation (`src/routes/AGENTS.md`) Complete implementation guide for future development: - Route function template and pattern - Required components (auth, session, logging, exceptions) - Anti-patterns and what NOT to do - Migration path from old code - Testing patterns with examples - Checklist for new routes ## 📊 Architecture Improvements ### Before ``` comfyui_client.py ├── HTTP transport (requests) ├── Response parsing ├── Polling logic ├── URL construction ├── Error handling ├── Retry logic └── Asset metadata extraction ``` ### After ``` src/ ├── auth/ # Authentication abstraction │ └── base.py # ComfyAuth, NoAuth, TokenAuth ├── client/ # HTTP client layer │ ├── get_data.py # Core async HTTP with httpx │ └── response.py # Standardized response object ├── utils/ # Cross-cutting concerns │ └── logging.py # @log_call decorator └── routes/ # Thin transport layer ├── workflow.py # Workflow operations ├── queue.py # Queue management ├── assets.py # Asset retrieval ├── models.py # Model discovery └── AGENTS.md # Implementation guide ``` ### Separation of Concerns | Layer | Responsibility | Example | |-------|----------------|---------| | **Routes** | HTTP transport only | `queue_workflow()` sends POST, returns response | | **Orchestrator** | Business logic, polling, validation | `ComfyUIClient` uses routes, handles retries | | **Managers** | Domain logic | `DefaultsManager`, `AssetRegistry` | | **Tools** | MCP tool interface | Tool registration, argument parsing | ## 🔄 Next Steps ### Phase 2: Refactor Orchestrator 1. Convert `comfyui_client.py` to use route functions 2. Move polling logic from routes to orchestrator 3. Add async/sync bridge (`asyncio.run()`) 4. Update managers to use routes where needed ### Phase 3: Testing 1. Add unit tests for route functions 2. Add integration tests for orchestrator 3. Add end-to-end tests for MCP tools ### Phase 4: Advanced Features 1. Add retry decorator (`@run_with_retry()`) 2. Implement response caching for GET requests 3. Add request context aggregation (`RouteContext`) 4. Consider full async conversion (if MCP supports it) ## 🎯 SOLID Principles Applied - ✅ **Single Responsibility**: Each route does ONE thing (HTTP transport) - ✅ **Open/Closed**: Easy to extend with new auth types without modifying routes - ✅ **Liskov Substitution**: All `ComfyAuth` subclasses are interchangeable - ✅ **Interface Segregation**: Minimal, focused interfaces per route - ✅ **Dependency Inversion**: Routes depend on `ComfyAuth` abstraction, not concrete classes ## 📦 Dependencies Updated Added to `pyproject.toml`: - `httpx>=0.27.0` - Modern async HTTP client for route functions - `fastmcp>=3.0.0` - FastMCP v3 for MCP server implementation - `pillow>=10.0.0` - Image processing (existing) Removed from direct dependencies: - `requests>=2.31.0` - Legacy sync HTTP (kept for backward compatibility in old code) - `mcp>=0.9.0` - Replaced by fastmcp v3 Now using **uv** for dependency management via `pyproject.toml`. ## 🧪 How to Use New Routes ```python import asyncio from src.auth import NoAuth from src.routes import queue_workflow # Create auth instance auth = NoAuth("http://127.0.0.1:8188") # Use route function async def example(): workflow = {"3": {"class_type": "KSampler", ...}} res = await queue_workflow(auth=auth, workflow=workflow) print(f"Queued: {res.response['prompt_id']}") # Run from sync code asyncio.run(example()) ``` ## 📝 Files Changed ### Created (16 new files) - `src/auth/__init__.py` - `src/auth/base.py` - `src/client/__init__.py` - `src/client/get_data.py` - `src/client/response.py` - `src/utils/__init__.py` - `src/utils/logging.py` - `src/routes/__init__.py` - `src/routes/workflow.py` - `src/routes/queue.py` - `src/routes/assets.py` - `src/routes/models.py` - `src/routes/AGENTS.md` - `REFACTORING_SUMMARY.md` (this file) ### Modified - `requirements.txt` (added httpx) ### To Be Refactored - `comfyui_client.py` → Use route functions - `server.py` → Update to use new auth module - `src/tools/*.py` → Update to use routes where applicable --- **Completed**: January 23, 2026 **Status**: Phase 1 Complete ✅