mcp-proxy-oauth-dcr 0.1.0__py3-none-any.whl

mcp_proxy/logging_config.py

@@ -0,0 +1,71 @@
+"""Logging configuration for the MCP Proxy.
+
+This module sets up structured logging using structlog for better
+debugging and monitoring capabilities.
+"""
+
+import logging
+import sys
+from typing import Any, Dict
+
+import structlog
+
+
+def configure_logging(log_level: str = "info") -> None:
+    """Configure structured logging for the application.
+    
+    Args:
+        log_level: Logging level (debug, info, warning, error, critical)
+    """
+    # Convert string level to logging constant
+    numeric_level = getattr(logging, log_level.upper(), logging.INFO)
+    
+    # Configure standard logging
+    logging.basicConfig(
+        format="%(message)s",
+        stream=sys.stderr,  # Use stderr to avoid interfering with stdio MCP
+        level=numeric_level,
+    )
+    
+    # Configure structlog
+    structlog.configure(
+        processors=[
+            structlog.contextvars.merge_contextvars,
+            structlog.processors.add_log_level,
+            structlog.processors.StackInfoRenderer(),
+            structlog.dev.set_exc_info,
+            structlog.processors.TimeStamper(fmt="iso", utc=True),
+            structlog.dev.ConsoleRenderer() if sys.stderr.isatty() else structlog.processors.JSONRenderer(),
+        ],
+        wrapper_class=structlog.make_filtering_bound_logger(numeric_level),
+        context_class=dict,
+        logger_factory=structlog.PrintLoggerFactory(file=sys.stderr),
+        cache_logger_on_first_use=True,
+    )
+
+
+def get_logger(name: str) -> structlog.BoundLogger:
+    """Get a logger instance for a module.
+    
+    Args:
+        name: Logger name (typically __name__)
+        
+    Returns:
+        Configured structlog logger
+    """
+    return structlog.get_logger(name)
+
+
+def add_context(**kwargs: Any) -> None:
+    """Add context variables to all subsequent log messages.
+    
+    Args:
+        **kwargs: Context key-value pairs
+    """
+    structlog.contextvars.clear_contextvars()
+    structlog.contextvars.bind_contextvars(**kwargs)
+
+
+def clear_context() -> None:
+    """Clear all context variables."""
+    structlog.contextvars.clear_contextvars()

mcp_proxy/models.py

@@ -0,0 +1,259 @@
+"""Core data models and interfaces for the MCP Proxy.
+
+This module defines the data structures used throughout the proxy system,
+including JSON-RPC messages, HTTP requests/responses, authentication state,
+and configuration models.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, HttpUrl, field_validator
+
+
+class MessageStatus(str, Enum):
+    """Status of message processing."""
+    PENDING = "pending"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+# ============================================================================
+# JSON-RPC Models
+# ============================================================================
+
+
+class JsonRpcError(BaseModel):
+    """JSON-RPC 2.0 error object."""
+    code: int = Field(..., description="Error code")
+    message: str = Field(..., description="Error message")
+    data: Optional[Any] = Field(None, description="Additional error data")
+
+
+class JsonRpcMessage(BaseModel):
+    """JSON-RPC 2.0 message structure."""
+    jsonrpc: str = Field(default="2.0", description="JSON-RPC version")
+    id: Optional[Union[str, int]] = Field(None, description="Request/response ID")
+    method: Optional[str] = Field(None, description="Method name for requests")
+    params: Optional[Any] = Field(None, description="Method parameters")
+    result: Optional[Any] = Field(None, description="Result for responses")
+    error: Optional[JsonRpcError] = Field(None, description="Error for error responses")
+
+    @field_validator("jsonrpc")
+    @classmethod
+    def validate_jsonrpc_version(cls, v: str) -> str:
+        """Ensure JSON-RPC version is 2.0."""
+        if v != "2.0":
+            raise ValueError("JSON-RPC version must be 2.0")
+        return v
+
+    def is_request(self) -> bool:
+        """Check if this is a request message."""
+        return self.method is not None
+
+    def is_response(self) -> bool:
+        """Check if this is a response message."""
+        return self.result is not None or self.error is not None
+
+    def is_notification(self) -> bool:
+        """Check if this is a notification (request without ID)."""
+        return self.method is not None and self.id is None
+
+    def to_json(self) -> str:
+        """Serialize to JSON string."""
+        return self.model_dump_json(exclude_none=True)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> JsonRpcMessage:
+        """Parse from JSON string."""
+        return cls.model_validate_json(json_str)
+
+
+# ============================================================================
+# HTTP MCP Models
+# ============================================================================
+
+
+class HttpMcpRequest(BaseModel):
+    """HTTP MCP request structure."""
+    method: str = Field(..., description="HTTP method (POST, GET)")
+    url: str = Field(..., description="Request URL")
+    headers: Dict[str, str] = Field(default_factory=dict, description="HTTP headers")
+    body: Optional[str] = Field(None, description="Request body")
+    session_id: Optional[str] = Field(None, description="MCP session ID")
+
+    @field_validator("method")
+    @classmethod
+    def validate_method(cls, v: str) -> str:
+        """Ensure HTTP method is valid."""
+        if v.upper() not in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
+            raise ValueError(f"Invalid HTTP method: {v}")
+        return v.upper()
+
+
+class HttpMcpResponse(BaseModel):
+    """HTTP MCP response structure."""
+    status: int = Field(..., description="HTTP status code")
+    headers: Dict[str, str] = Field(default_factory=dict, description="Response headers")
+    body: str = Field(default="", description="Response body")
+    content_type: str = Field(default="application/json", description="Content type")
+
+    def is_success(self) -> bool:
+        """Check if response indicates success."""
+        return 200 <= self.status < 300
+
+    def is_auth_error(self) -> bool:
+        """Check if response indicates authentication error."""
+        return self.status == 401
+
+
+# ============================================================================
+# Authentication Models
+# ============================================================================
+
+
+class ClientCredentials(BaseModel):
+    """OAuth client credentials."""
+    client_id: str = Field(..., description="OAuth client ID")
+    client_secret: str = Field(..., description="OAuth client secret")
+    expires_at: Optional[datetime] = Field(None, description="Credential expiration time")
+
+    def is_expired(self) -> bool:
+        """Check if credentials are expired."""
+        if self.expires_at is None:
+            return False
+        return datetime.now() >= self.expires_at
+
+
+class OAuthTokenResponse(BaseModel):
+    """OAuth token response."""
+    access_token: str = Field(..., description="Access token")
+    token_type: str = Field(..., description="Token type (usually Bearer)")
+    expires_in: int = Field(..., description="Token lifetime in seconds")
+    scope: Optional[str] = Field(None, description="Token scope")
+    refresh_token: Optional[str] = Field(None, description="Refresh token")
+
+    def get_expiration_time(self) -> datetime:
+        """Calculate token expiration time."""
+        from datetime import timedelta
+        return datetime.now() + timedelta(seconds=self.expires_in)
+
+
+class AuthenticationState(BaseModel):
+    """Current authentication state."""
+    client_credentials: Optional[ClientCredentials] = Field(
+        None, description="Client credentials"
+    )
+    access_token: Optional[str] = Field(None, description="Current access token")
+    refresh_token: Optional[str] = Field(None, description="Refresh token for token renewal")
+    token_expires_at: Optional[datetime] = Field(
+        None, description="Token expiration time"
+    )
+    last_dcr_attempt: Optional[datetime] = Field(
+        None, description="Last DCR attempt timestamp"
+    )
+    dcr_retry_count: int = Field(default=0, description="DCR retry counter")
+    is_authenticated: bool = Field(default=False, description="Authentication status")
+
+    def is_token_valid(self) -> bool:
+        """Check if current token is valid."""
+        if not self.access_token or not self.token_expires_at:
+            return False
+        # Consider token invalid 60 seconds before actual expiration
+        from datetime import timedelta
+        return datetime.now() < (self.token_expires_at - timedelta(seconds=60))
+
+    def needs_token_refresh(self) -> bool:
+        """Check if token needs refresh."""
+        return self.is_authenticated and not self.is_token_valid()
+
+
+# ============================================================================
+# Session Models
+# ============================================================================
+
+
+class SessionState(BaseModel):
+    """MCP session state."""
+    session_id: Optional[str] = Field(None, description="Session ID")
+    is_initialized: bool = Field(default=False, description="Initialization status")
+    protocol_version: str = Field(default="1.0", description="MCP protocol version")
+    capabilities: List[str] = Field(default_factory=list, description="Session capabilities")
+    sse_stream_active: bool = Field(default=False, description="SSE stream status")
+
+
+# ============================================================================
+# Message Correlation Models
+# ============================================================================
+
+
+@dataclass
+class MessageCorrelation:
+    """Tracks correlation between stdio and HTTP messages."""
+    stdio_request_id: Union[str, int]
+    http_request_id: str
+    timestamp: datetime
+    method: str
+    status: MessageStatus = MessageStatus.PENDING
+
+    def is_pending(self) -> bool:
+        """Check if message is still pending."""
+        return self.status == MessageStatus.PENDING
+
+    def mark_completed(self) -> None:
+        """Mark message as completed."""
+        self.status = MessageStatus.COMPLETED
+
+    def mark_failed(self) -> None:
+        """Mark message as failed."""
+        self.status = MessageStatus.FAILED
+
+
+# ============================================================================
+# Configuration Models
+# ============================================================================
+
+
+class ProxyConfig(BaseModel):
+    """Proxy configuration."""
+    mcp_server_url: HttpUrl = Field(..., description="HTTP MCP server URL")
+    oauth_provider_url: HttpUrl = Field(..., description="OAuth provider URL")
+    client_name: str = Field(default="mcp-proxy-client", description="OAuth client name")
+    scopes: List[str] = Field(
+        default_factory=lambda: ["mcp:read", "mcp:write"],
+        description="OAuth scopes"
+    )
+    connection_timeout: int = Field(
+        default=30, ge=1, le=300, description="Connection timeout in seconds"
+    )
+    retry_attempts: int = Field(
+        default=3, ge=0, le=10, description="Number of retry attempts"
+    )
+    log_level: str = Field(
+        default="info", description="Logging level"
+    )
+    max_backoff_seconds: int = Field(
+        default=60, ge=1, le=300, description="Maximum backoff time in seconds"
+    )
+
+    @field_validator("log_level")
+    @classmethod
+    def validate_log_level(cls, v: str) -> str:
+        """Ensure log level is valid."""
+        valid_levels = ["debug", "info", "warning", "error", "critical"]
+        if v.lower() not in valid_levels:
+            raise ValueError(f"Invalid log level: {v}. Must be one of {valid_levels}")
+        return v.lower()
+
+    @field_validator("scopes")
+    @classmethod
+    def validate_scopes(cls, v: List[str]) -> List[str]:
+        """Ensure scopes list is not empty."""
+        if not v:
+            raise ValueError("Scopes list cannot be empty")
+        return v

mcp_proxy/protocols.py

@@ -0,0 +1,122 @@
+"""Protocol interfaces for the MCP Proxy components.
+
+This module defines the abstract interfaces (protocols) that components
+must implement, enabling loose coupling and testability.
+"""
+
+from typing import Protocol, Optional, AsyncIterator
+from .models import (
+    JsonRpcMessage,
+    HttpMcpRequest,
+    HttpMcpResponse,
+    ClientCredentials,
+    OAuthTokenResponse,
+    ProxyConfig,
+    AuthenticationState,
+)
+
+
+class StdioInterface(Protocol):
+    """Interface for stdio communication with Kiro."""
+
+    async def start(self) -> None:
+        """Start the stdio interface."""
+        ...
+
+    async def stop(self) -> None:
+        """Stop the stdio interface."""
+        ...
+
+    async def send_message(self, message: JsonRpcMessage) -> None:
+        """Send a JSON-RPC message to stdout."""
+        ...
+
+    async def receive_message(self) -> JsonRpcMessage:
+        """Receive a JSON-RPC message from stdin."""
+        ...
+
+
+class ProtocolTranslator(Protocol):
+    """Interface for protocol translation between stdio and HTTP."""
+
+    async def translate_stdio_to_http(
+        self, message: JsonRpcMessage
+    ) -> HttpMcpRequest:
+        """Translate stdio JSON-RPC message to HTTP MCP request."""
+        ...
+
+    async def translate_http_to_stdio(
+        self, response: HttpMcpResponse
+    ) -> JsonRpcMessage:
+        """Translate HTTP MCP response to stdio JSON-RPC message."""
+        ...
+
+    def correlate_messages(self, request_id: str, response_id: str) -> None:
+        """Correlate request and response messages."""
+        ...
+
+
+class AuthenticationManager(Protocol):
+    """Interface for OAuth DCR and token management."""
+
+    async def initialize(self) -> None:
+        """Initialize authentication manager and perform DCR if needed."""
+        ...
+
+    async def get_access_token(self) -> str:
+        """Get a valid access token, refreshing if necessary."""
+        ...
+
+    async def refresh_token(self) -> str:
+        """Refresh the access token."""
+        ...
+
+    async def perform_dcr(self) -> ClientCredentials:
+        """Perform OAuth Dynamic Client Registration."""
+        ...
+
+    def is_token_valid(self) -> bool:
+        """Check if current token is valid."""
+        ...
+
+    def get_state(self) -> AuthenticationState:
+        """Get current authentication state."""
+        ...
+
+
+class HttpClient(Protocol):
+    """Interface for HTTP communication with backend MCP server."""
+
+    async def send_request(self, request: HttpMcpRequest) -> HttpMcpResponse:
+        """Send HTTP request to backend MCP server."""
+        ...
+
+    async def open_sse_stream(
+        self, session_id: Optional[str] = None
+    ) -> AsyncIterator[str]:
+        """Open Server-Sent Events stream."""
+        ...
+
+    async def close_connection(self) -> None:
+        """Close HTTP connection."""
+        ...
+
+    def set_auth_token(self, token: str) -> None:
+        """Set authentication token for requests."""
+        ...
+
+
+class ConfigurationManager(Protocol):
+    """Interface for configuration management."""
+
+    async def load(self) -> ProxyConfig:
+        """Load configuration from environment and files."""
+        ...
+
+    def validate(self, config: ProxyConfig) -> bool:
+        """Validate configuration parameters."""
+        ...
+
+    def get_defaults(self) -> ProxyConfig:
+        """Get default configuration values."""
+        ...