PR feedback updates

Author: peteski22

.gitignore

@@ -48,9 +48,6 @@ dmypy.json
 # Ruff
 .ruff_cache/
 
-# Distribution / packaging
-*.egg-info/
-
 # Proto source files (downloaded during build, not committed)
 tmp/
 

Makefile

@@ -1,22 +1,23 @@
-.PHONY: help
+.PHONY: all help
+all: help ## Default target
 help: ## Show this help message
 	@echo "Available targets:"
 	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "  \033[36m%-20s\033[0m %s\n", $$1, $$2}'
 
 .PHONY: ensure-scripts-exec
 ensure-scripts-exec: ## Make scripts executable
-	chmod +x scripts/*
+	@if [ -d scripts ]; then chmod +x scripts/*.sh 2>/dev/null || true; fi
 
 .PHONY: setup
 setup: ensure-scripts-exec ## Setup development environment (installs uv and syncs dependencies)
 	./scripts/setup_uv.sh
 
 .PHONY: test
-test: ## Run tests with pytest
+test: setup ## Run tests with pytest
 	uv run pytest tests/ -v
 
 .PHONY: lint
-lint: ## Run pre-commit hooks on all files
+lint: setup ## Run pre-commit hooks on all files
 	uv run pre-commit run --all-files
 
 .PHONY: generate-protos
@@ -39,12 +40,20 @@ build-plugin-prod: ensure-scripts-exec ## Build a plugin with Nuitka for product
 	fi
 	./scripts/build_plugin.sh $(PLUGIN) --nuitka
 
-.PHONY: clean
-clean: ## Clean generated files and caches
-	rm -rf tmp/
-	find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
-	find . -type f -name "*.pyc" -delete
+.PHONY: clean clean-build clean-caches clean-pyc
+clean: clean-build clean-caches clean-pyc ## Clean generated files and caches
+
+.PHONY: clean-build
+clean-build: ## Clean build artifacts
+	rm -rf build/ dist/ tmp/
+	find . -type d -name "*.egg-info" -exec rm -rf {} + 2>/dev/null || true
+
+.PHONY: clean-caches
+clean-caches: ## Clean cache directories
 	find . -type d -name ".pytest_cache" -exec rm -rf {} + 2>/dev/null || true
 	find . -type d -name ".ruff_cache" -exec rm -rf {} + 2>/dev/null || true
-	find . -type d -name "*.egg-info" -exec rm -rf {} + 2>/dev/null || true
-	rm -rf build/ dist/
+	find . -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true
+
+.PHONY: clean-pyc
+clean-pyc: ## Clean Python bytecode files
+	find . -type f -name "*.pyc" -delete

README.md

@@ -116,6 +116,7 @@ When handling requests or responses, you can:
 The SDK includes five example plugins demonstrating common patterns:
 
 ### 1. Simple Plugin
+
 Adds a custom header to all requests.
 
 ```bash
@@ -124,6 +125,7 @@ uv run python main.py
 ```
 
 ### 2. Auth Plugin
+
 Validates Bearer token authentication and rejects unauthorized requests.
 
 ```bash
@@ -133,6 +135,7 @@ uv run python main.py
 ```
 
 ### 3. Logging Plugin
+
 Logs HTTP request and response details for observability.
 
 ```bash
@@ -141,6 +144,7 @@ uv run python main.py
 ```
 
 ### 4. Rate Limit Plugin
+
 Implements token bucket rate limiting per client IP.
 
 ```bash
@@ -149,6 +153,7 @@ uv run python main.py
 ```
 
 ### 5. Transform Plugin
+
 Transforms JSON request bodies by adding metadata fields.
 
 ```bash
@@ -214,24 +219,22 @@ class BasePlugin(PluginServicer):
     async def CheckHealth(self, request: Empty, context) -> Empty
     async def CheckReady(self, request: Empty, context) -> Empty
     async def HandleRequest(self, request: HTTPRequest, context) -> HTTPResponse
-    async def HandleResponse(self, request: HTTPResponse, context) -> HTTPResponse
+    async def HandleResponse(self, response: HTTPResponse, context) -> HTTPResponse
 ```
 
-### serve()
+### `serve()`
 
 ```python
 async def serve(
     plugin: BasePlugin,
     args: Optional[list[str]] = None,  # Command-line arguments (typically sys.argv)
-    max_workers: int = 10,
     grace_period: float = 5.0,
 ) -> None
 ```
 
 **Parameters:**
 - `plugin`: The plugin instance to serve
 - `args`: Command-line arguments. When provided (e.g., `sys.argv`), enables mcpd compatibility by parsing `--address` and `--network` flags. When `None`, runs in standalone mode on TCP port 50051.
-- `max_workers`: Maximum number of concurrent gRPC workers
 - `grace_period`: Seconds to wait during graceful shutdown
 
 **Command-line flags** (when `args` is provided):

examples/auth_plugin/__init__.py

@@ -1 +1,5 @@
 """Authentication plugin example."""
+
+from .main import AuthPlugin
+
+__all__ = ["AuthPlugin"]

examples/auth_plugin/main.py

@@ -9,7 +9,7 @@
 import sys
 
 from google.protobuf.empty_pb2 import Empty
-from grpc import ServicerContext
+from grpc.aio import ServicerContext
 
 from mcpd_plugins import BasePlugin, serve
 from mcpd_plugins.v1.plugins.plugin_pb2 import (
@@ -23,6 +23,9 @@
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Authentication scheme.
+BEARER_SCHEME = "Bearer"
+
 
 class AuthPlugin(BasePlugin):
     """Plugin that validates Bearer token authentication."""
@@ -46,17 +49,17 @@ async def GetCapabilities(self, request: Empty, context: ServicerContext) -> Cap
 
     async def HandleRequest(self, request: HTTPRequest, context: ServicerContext) -> HTTPResponse:
         """Validate Bearer token in Authorization header."""
-        logger.info(f"Authenticating request: {request.method} {request.url}")
+        logger.info("Authenticating request: %s %s", request.method, request.url)
 
         # Check for Authorization header.
         auth_header = request.headers.get("Authorization", "")
 
-        if not auth_header.startswith("Bearer "):
+        if not auth_header.startswith(f"{BEARER_SCHEME} "):
             logger.warning("Missing or invalid Authorization header")
             return self._unauthorized_response("Missing or invalid Authorization header")
 
         # Extract and validate token.
-        token = auth_header[7:]  # Remove "Bearer " prefix
+        token = auth_header.removeprefix(f"{BEARER_SCHEME} ")
         if token != self.expected_token:
             logger.warning("Invalid token")
             return self._unauthorized_response("Invalid token")
@@ -73,7 +76,7 @@ def _unauthorized_response(self, message: str) -> HTTPResponse:
             **{"continue": False},
         )
         response.headers["Content-Type"] = "application/json"
-        response.headers["WWW-Authenticate"] = "Bearer"
+        response.headers["WWW-Authenticate"] = BEARER_SCHEME
         return response
 
 

examples/logging_plugin/__init__.py

@@ -1 +1,5 @@
 """Logging plugin example."""
+
+from .main import LoggingPlugin
+
+__all__ = ["LoggingPlugin"]

examples/logging_plugin/main.py

@@ -9,7 +9,7 @@
 import sys
 
 from google.protobuf.empty_pb2 import Empty
-from grpc import ServicerContext
+from grpc.aio import ServicerContext
 
 from mcpd_plugins import BasePlugin, serve
 from mcpd_plugins.v1.plugins.plugin_pb2 import (
@@ -47,42 +47,42 @@ async def HandleRequest(self, request: HTTPRequest, context: ServicerContext) ->
         """Log incoming request details."""
         logger.info("=" * 80)
         logger.info("INCOMING REQUEST")
-        logger.info(f"Method: {request.method}")
-        logger.info(f"URL: {request.url}")
-        logger.info(f"Path: {request.path}")
-        logger.info(f"Remote Address: {request.remote_addr}")
+        logger.info("Method: %s", request.method)
+        logger.info("URL: %s", request.url)
+        logger.info("Path: %s", request.path)
+        logger.info("Remote Address: %s", request.remote_addr)
 
         # Log headers.
         logger.info("Headers:")
         for key, value in request.headers.items():
             # Mask sensitive headers.
             if key.lower() in ("authorization", "cookie"):
                 value = "***REDACTED***"
-            logger.info(f"  {key}: {value}")
+            logger.info("  %s: %s", key, value)
 
         # Log body size.
         if request.body:
-            logger.info(f"Body size: {len(request.body)} bytes")
+            logger.info("Body size: %s bytes", len(request.body))
 
         logger.info("=" * 80)
 
         # Continue processing.
         return HTTPResponse(**{"continue": True})
 
-    async def HandleResponse(self, request: HTTPResponse, context: ServicerContext) -> HTTPResponse:
+    async def HandleResponse(self, response: HTTPResponse, context: ServicerContext) -> HTTPResponse:
         """Log outgoing response details."""
         logger.info("=" * 80)
         logger.info("OUTGOING RESPONSE")
-        logger.info(f"Status Code: {request.status_code}")
+        logger.info("Status Code: %s", response.status_code)
 
         # Log headers.
         logger.info("Headers:")
-        for key, value in request.headers.items():
-            logger.info(f"  {key}: {value}")
+        for key, value in response.headers.items():
+            logger.info("  %s: %s", key, value)
 
         # Log body size.
-        if request.body:
-            logger.info(f"Body size: {len(request.body)} bytes")
+        if response.body:
+            logger.info("Body size: %s bytes", len(response.body))
 
         logger.info("=" * 80)
 

examples/rate_limit_plugin/__init__.py

@@ -1 +1,5 @@
 """Rate limiting plugin example."""
+
+from .main import RateLimitPlugin
+
+__all__ = ["RateLimitPlugin"]

examples/rate_limit_plugin/main.py

@@ -10,7 +10,7 @@
 from collections import defaultdict
 
 from google.protobuf.empty_pb2 import Empty
-from grpc import ServicerContext
+from grpc.aio import ServicerContext
 
 from mcpd_plugins import BasePlugin, serve
 from mcpd_plugins.v1.plugins.plugin_pb2 import (
@@ -41,6 +41,7 @@ def __init__(self, requests_per_minute: int = 60):
         # Track tokens for each client IP.
         self.buckets: dict[str, float] = defaultdict(lambda: float(requests_per_minute))
         self.last_update: dict[str, float] = defaultdict(time.time)
+        self.locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
 
     async def GetMetadata(self, request: Empty, context: ServicerContext) -> Metadata:
         """Return plugin metadata."""
@@ -57,33 +58,34 @@ async def GetCapabilities(self, request: Empty, context: ServicerContext) -> Cap
     async def HandleRequest(self, request: HTTPRequest, context: ServicerContext) -> HTTPResponse:
         """Apply rate limiting based on client IP."""
         client_ip = request.remote_addr or "unknown"
-        logger.info(f"Rate limit check for {client_ip}: {request.method} {request.url}")
-
-        # Refill tokens based on time elapsed.
-        now = time.time()
-        elapsed = now - self.last_update[client_ip]
-        self.buckets[client_ip] = min(
-            self.requests_per_minute,
-            self.buckets[client_ip] + elapsed * self.rate_per_second,
-        )
-        self.last_update[client_ip] = now
-
-        # Check if client has tokens available.
-        if self.buckets[client_ip] < 1.0:
-            logger.warning(f"Rate limit exceeded for {client_ip}")
-            return self._rate_limit_response(client_ip)
-
-        # Consume one token.
-        self.buckets[client_ip] -= 1.0
-        logger.info(f"Request allowed for {client_ip} (tokens remaining: {self.buckets[client_ip]:.2f})")
-
-        # Add rate limit headers to response.
-        response = HTTPResponse(**{"continue": True})
-        response.modified_request.CopyFrom(request)
-        response.headers["X-RateLimit-Limit"] = str(self.requests_per_minute)
-        response.headers["X-RateLimit-Remaining"] = str(int(self.buckets[client_ip]))
-
-        return response
+        logger.info("Rate limit check for %s: %s %s", client_ip, request.method, request.url)
+
+        async with self.locks[client_ip]:
+            # Refill tokens based on time elapsed.
+            now = time.time()
+            elapsed = now - self.last_update[client_ip]
+            self.buckets[client_ip] = min(
+                self.requests_per_minute,
+                self.buckets[client_ip] + elapsed * self.rate_per_second,
+            )
+            self.last_update[client_ip] = now
+
+            # Check if client has tokens available.
+            if self.buckets[client_ip] < 1.0:
+                logger.warning("Rate limit exceeded for %s", client_ip)
+                return self._rate_limit_response(client_ip)
+
+            # Consume one token.
+            self.buckets[client_ip] -= 1.0
+            logger.info("Request allowed for %s (tokens remaining: %.2f)", client_ip, self.buckets[client_ip])
+
+            # Add rate limit headers to response.
+            response = HTTPResponse(**{"continue": True})
+            response.modified_request.CopyFrom(request)
+            response.headers["X-RateLimit-Limit"] = str(self.requests_per_minute)
+            response.headers["X-RateLimit-Remaining"] = str(int(self.buckets[client_ip]))
+
+            return response
 
     def _rate_limit_response(self, client_ip: str) -> HTTPResponse:
         """Create a 429 Too Many Requests response."""

examples/simple_plugin/__init__.py

@@ -1 +1,5 @@
 """Simple plugin example."""
+
+from .main import SimplePlugin
+
+__all__ = ["SimplePlugin"]