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

mcp_proxy/translator/translator.py

@@ -0,0 +1,691 @@
+"""Protocol translator implementation for MCP Proxy.
+
+This module provides translation between stdio JSON-RPC messages and HTTP MCP requests,
+handling method mapping, request correlation, and protocol-specific differences with
+comprehensive error logging and diagnostics.
+"""
+
+import json
+import logging
+from typing import Dict, Optional, Union
+from datetime import datetime
+
+from ..models import (
+    JsonRpcMessage,
+    JsonRpcError,
+    HttpMcpRequest,
+    HttpMcpResponse,
+    MessageCorrelation,
+    MessageStatus,
+)
+from ..stdio.jsonrpc import MessageIdGenerator, MessageCorrelationTracker
+from ..exceptions import (
+    ProtocolError,
+    UnsupportedMethodError,
+    InvalidJsonRpcError,
+    MessageCorrelationError,
+)
+from ..logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class ProtocolTranslatorImpl:
+    """Translates between stdio JSON-RPC and HTTP MCP protocols.
+    
+    This translator handles:
+    - Converting JSON-RPC messages to HTTP POST requests
+    - Mapping MCP methods between protocols
+    - Maintaining request correlation across protocols
+    - Managing session state and headers
+    
+    Attributes:
+        base_url: Base URL for the HTTP MCP server
+        session_id: Current MCP session ID
+        correlation_tracker: Tracks message correlations
+        id_generator: Generates unique HTTP request IDs
+    """
+    
+    # MCP HTTP Streamable transport sends ALL requests to the same endpoint (base URL)
+    # The method name is included in the JSON-RPC body, not the URL path
+    # This is different from REST-style APIs where each method has its own endpoint
+    # Use empty string "" instead of "/" to avoid trailing slash issues
+    METHOD_ENDPOINT_MAP = {
+        "initialize": "",
+        "ping": "",
+        "tools/list": "",
+        "tools/call": "",
+        "resources/list": "",
+        "resources/read": "",
+        "prompts/list": "",
+        "prompts/get": "",
+        "completion/complete": "",
+        "logging/setLevel": "",
+        "notifications/initialized": "",
+        "notifications/cancelled": "",
+        "notifications/progress": "",
+    }
+    
+    # HTTP method for each MCP method (all are POST for JSON-RPC over HTTP)
+    METHOD_HTTP_METHOD_MAP = {
+        "initialize": "POST",
+        "ping": "POST",
+        "tools/list": "POST",
+        "tools/call": "POST",
+        "resources/list": "POST",
+        "resources/read": "POST",
+        "prompts/list": "POST",
+        "prompts/get": "POST",
+        "completion/complete": "POST",
+        "logging/setLevel": "POST",
+        "notifications/initialized": "POST",
+        "notifications/cancelled": "POST",
+        "notifications/progress": "POST",
+    }
+    
+    def __init__(self, base_url: str, session_id: Optional[str] = None):
+        """Initialize the protocol translator.
+        
+        Args:
+            base_url: Base URL for the HTTP MCP server
+            session_id: Optional initial session ID
+        """
+        self._base_url = base_url.rstrip("/")
+        self._session_id = session_id
+        self._correlation_tracker = MessageCorrelationTracker()
+        self._id_generator = MessageIdGenerator(prefix="http")
+        
+        logger.debug(
+            "ProtocolTranslatorImpl initialized with base_url=%s, session_id=%s",
+            self._base_url, self._session_id
+        )
+    
+    @property
+    def session_id(self) -> Optional[str]:
+        """Get the current session ID."""
+        return self._session_id
+    
+    @session_id.setter
+    def session_id(self, value: Optional[str]) -> None:
+        """Set the session ID."""
+        self._session_id = value
+        logger.debug("Session ID updated to: %s", value)
+    
+    @property
+    def correlation_tracker(self) -> MessageCorrelationTracker:
+        """Get the correlation tracker."""
+        return self._correlation_tracker
+    
+    def _extract_json_from_sse(self, sse_body: str) -> str:
+        """Extract JSON data from SSE format response.
+        
+        SSE format looks like:
+            event: message
+            data: {"jsonrpc":"2.0",...}
+        
+        This method extracts the JSON from the "data:" line(s).
+        
+        Args:
+            sse_body: Raw SSE response body
+            
+        Returns:
+            Extracted JSON string, or empty string if no data found
+        """
+        if not sse_body:
+            return ""
+        
+        data_lines = []
+        for line in sse_body.split("\n"):
+            line = line.strip()
+            if line.startswith("data:"):
+                # Extract content after "data:"
+                data_content = line[5:].strip()
+                if data_content:
+                    data_lines.append(data_content)
+        
+        # Join multiple data lines (SSE spec allows multi-line data)
+        return "\n".join(data_lines) if data_lines else ""
+    
+    def _get_endpoint_for_method(self, method: str) -> str:
+        """Get the HTTP endpoint for an MCP method.
+        
+        For MCP HTTP Streamable transport, all methods go to the same endpoint (base URL).
+        The method name is included in the JSON-RPC body.
+        
+        Args:
+            method: MCP method name
+            
+        Returns:
+            HTTP endpoint path (always "" for MCP HTTP Streamable to avoid trailing slash)
+        """
+        # MCP HTTP Streamable sends all requests to the base URL
+        # Return "" for all methods - the method is in the JSON-RPC body
+        return self.METHOD_ENDPOINT_MAP.get(method, "")
+    
+    def _get_http_method_for_mcp_method(self, method: str) -> str:
+        """Get the HTTP method for an MCP method.
+        
+        For MCP HTTP Streamable transport, all requests use POST.
+        
+        Args:
+            method: MCP method name
+            
+        Returns:
+            HTTP method (always POST for MCP)
+        """
+        # MCP HTTP Streamable uses POST for all JSON-RPC requests
+        return self.METHOD_HTTP_METHOD_MAP.get(method, "POST")
+    
+    def _build_http_body(self, message: JsonRpcMessage) -> Optional[str]:
+        """Build HTTP request body from JSON-RPC message.
+        
+        For POST requests, the body contains the JSON-RPC message parameters.
+        
+        Args:
+            message: JSON-RPC message
+            
+        Returns:
+            JSON string for request body, or None for GET requests
+        """
+        if not message.is_request():
+            raise InvalidJsonRpcError(
+                "Cannot build HTTP body from non-request message",
+                details={"message": message.model_dump()}
+            )
+        
+        # Build request body with params
+        body_data = {
+            "jsonrpc": "2.0",
+            "id": message.id,
+            "method": message.method,
+        }
+        
+        # Include params if present
+        if message.params is not None:
+            body_data["params"] = message.params
+        
+        return json.dumps(body_data)
+    
+    async def translate_stdio_to_http(
+        self,
+        message: JsonRpcMessage
+    ) -> HttpMcpRequest:
+        """Translate stdio JSON-RPC message to HTTP MCP request.
+        
+        This method:
+        1. Validates the JSON-RPC message is a request
+        2. Maps the MCP method to an HTTP endpoint
+        3. Generates a unique HTTP request ID
+        4. Creates correlation between stdio and HTTP IDs
+        5. Builds the HTTP request with appropriate headers and body
+        
+        Args:
+            message: JSON-RPC message from stdio
+            
+        Returns:
+            HTTP MCP request ready to send
+            
+        Raises:
+            InvalidJsonRpcError: If message is not a valid request
+            UnsupportedMethodError: If method is not supported
+            MessageCorrelationError: If correlation already exists
+        """
+        # Validate this is a request message
+        if not message.is_request():
+            raise InvalidJsonRpcError(
+                "Can only translate request messages to HTTP",
+                details={"message": message.model_dump()}
+            )
+        
+        method = message.method
+        if not method:
+            raise InvalidJsonRpcError(
+                "Request message must have a method",
+                details={"message": message.model_dump()}
+            )
+        
+        logger.debug(
+            "Translating stdio request to HTTP: method=%s, id=%s",
+            method, message.id
+        )
+        
+        # Get endpoint and HTTP method for this MCP method
+        endpoint = self._get_endpoint_for_method(method)
+        http_method = self._get_http_method_for_mcp_method(method)
+        
+        # Generate unique HTTP request ID
+        http_request_id = self._id_generator.generate()
+        
+        # Create correlation if this is not a notification
+        if message.id is not None:
+            self._correlation_tracker.add_correlation(
+                stdio_request_id=message.id,
+                http_request_id=http_request_id,
+                method=method
+            )
+            logger.debug(
+                "Created correlation: stdio_id=%s -> http_id=%s",
+                message.id, http_request_id
+            )
+        
+        # Build request body
+        body = self._build_http_body(message)
+        
+        # Build headers - MCP HTTP Streamable requires Accept header for both JSON and SSE
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+        }
+        
+        # Add session ID header if available
+        if self._session_id:
+            headers["Mcp-Session-Id"] = self._session_id
+        
+        # Create HTTP request
+        http_request = HttpMcpRequest(
+            method=http_method,
+            url=f"{self._base_url}{endpoint}",
+            headers=headers,
+            body=body,
+            session_id=self._session_id,
+        )
+        
+        logger.info(
+            "Translated stdio request to HTTP: method=%s, endpoint=%s, "
+            "stdio_id=%s, http_id=%s",
+            method, endpoint, message.id, http_request_id
+        )
+        
+        return http_request
+    
+    async def translate_http_to_stdio(
+        self,
+        response: HttpMcpResponse,
+        stdio_request_id: Optional[Union[str, int]] = None
+    ) -> JsonRpcMessage:
+        """Translate HTTP MCP response to stdio JSON-RPC message.
+        
+        This method:
+        1. Parses the HTTP response body as JSON (or extracts from SSE format)
+        2. Handles successful responses by extracting the result
+        3. Handles error responses by translating to JSON-RPC errors
+        4. Maintains correlation with the original stdio request
+        
+        Args:
+            response: HTTP MCP response
+            stdio_request_id: Optional stdio request ID for correlation lookup
+            
+        Returns:
+            JSON-RPC message for stdio
+            
+        Raises:
+            ProtocolError: If response cannot be translated
+            MessageCorrelationError: If correlation cannot be found
+        """
+        logger.debug(
+            "Translating HTTP response to stdio: status=%s, content_type=%s, stdio_id=%s",
+            response.status, response.content_type, stdio_request_id
+        )
+        
+        # If stdio_request_id not provided, we can't correlate
+        # This is acceptable for notifications or streaming events
+        request_id = stdio_request_id
+        
+        # Handle HTTP errors by translating to JSON-RPC errors
+        if not response.is_success():
+            return self._translate_http_error_to_stdio(response, request_id)
+        
+        # Extract body data - handle SSE format if content type is text/event-stream
+        body_to_parse = response.body
+        if response.content_type and "text/event-stream" in response.content_type:
+            # SSE format: extract JSON from "data:" lines
+            body_to_parse = self._extract_json_from_sse(response.body)
+        
+        # Parse response body
+        try:
+            body_data = json.loads(body_to_parse) if body_to_parse else {}
+        except json.JSONDecodeError as e:
+            logger.error("Failed to parse HTTP response body: %s", e)
+            # Return JSON-RPC parse error
+            return JsonRpcMessage(
+                jsonrpc="2.0",
+                id=request_id,
+                error=JsonRpcError(
+                    code=-32700,
+                    message="Parse error",
+                    data={"details": f"Invalid JSON in HTTP response: {str(e)}"}
+                )
+            )
+        
+        # Check if the response body is already a JSON-RPC message
+        if "jsonrpc" in body_data:
+            # Response is already in JSON-RPC format, use it directly
+            # but ensure the ID matches our correlation
+            if request_id is not None:
+                body_data["id"] = request_id
+            
+            try:
+                stdio_message = JsonRpcMessage.model_validate(body_data)
+                logger.info(
+                    "Translated HTTP response to stdio: status=%s, stdio_id=%s",
+                    response.status, request_id
+                )
+                return stdio_message
+            except Exception as e:
+                logger.error("Failed to validate JSON-RPC message: %s", e)
+                return JsonRpcMessage(
+                    jsonrpc="2.0",
+                    id=request_id,
+                    error=JsonRpcError(
+                        code=-32603,
+                        message="Internal error",
+                        data={"details": f"Invalid JSON-RPC format: {str(e)}"}
+                    )
+                )
+        
+        # Response is not in JSON-RPC format, wrap it as a result
+        stdio_message = JsonRpcMessage(
+            jsonrpc="2.0",
+            id=request_id,
+            result=body_data
+        )
+        
+        logger.info(
+            "Translated HTTP response to stdio: status=%s, stdio_id=%s",
+            response.status, request_id
+        )
+        
+        return stdio_message
+    
+    def _translate_http_error_to_stdio(
+        self,
+        response: HttpMcpResponse,
+        request_id: Optional[Union[str, int]] = None
+    ) -> JsonRpcMessage:
+        """Translate HTTP error response to JSON-RPC error message.
+        
+        Maps HTTP status codes to appropriate JSON-RPC error codes:
+        - 400 Bad Request -> -32602 Invalid params
+        - 401 Unauthorized -> -32000 Server error (authentication)
+        - 404 Not Found -> -32601 Method not found
+        - 500+ Server errors -> -32603 Internal error
+        - Other errors -> -32603 Internal error
+        
+        Args:
+            response: HTTP error response
+            request_id: Optional stdio request ID
+            
+        Returns:
+            JSON-RPC error message
+        """
+        logger.debug(
+            "Translating HTTP error to stdio: status=%s, stdio_id=%s",
+            response.status, request_id
+        )
+        
+        # Try to parse error details from response body
+        error_data = None
+        error_message = f"HTTP {response.status}"
+        
+        try:
+            if response.body:
+                body_data = json.loads(response.body)
+                
+                # Check if body contains a JSON-RPC error
+                if "error" in body_data:
+                    error_obj = body_data["error"]
+                    if isinstance(error_obj, dict):
+                        # Use the JSON-RPC error directly
+                        return JsonRpcMessage(
+                            jsonrpc="2.0",
+                            id=request_id,
+                            error=JsonRpcError(
+                                code=error_obj.get("code", -32603),
+                                message=error_obj.get("message", error_message),
+                                data=error_obj.get("data")
+                            )
+                        )
+                
+                # Extract error message from common formats
+                if "message" in body_data:
+                    error_message = body_data["message"]
+                elif "error" in body_data and isinstance(body_data["error"], str):
+                    error_message = body_data["error"]
+                
+                # Store full body as error data
+                error_data = body_data
+        except json.JSONDecodeError:
+            # If body is not JSON, use it as error message
+            if response.body:
+                error_message = response.body[:200]  # Limit length
+                error_data = {"raw_response": response.body[:500]}
+        
+        # Map HTTP status to JSON-RPC error code
+        if response.status == 400:
+            error_code = -32602  # Invalid params
+        elif response.status == 401:
+            error_code = -32000  # Server error (authentication)
+            error_message = f"Authentication failed: {error_message}"
+        elif response.status == 404:
+            error_code = -32601  # Method not found
+            error_message = f"Method not found: {error_message}"
+        elif response.status >= 500:
+            error_code = -32603  # Internal error
+            error_message = f"Server error: {error_message}"
+        else:
+            error_code = -32603  # Internal error
+            error_message = f"HTTP error {response.status}: {error_message}"
+        
+        # Create JSON-RPC error message
+        stdio_message = JsonRpcMessage(
+            jsonrpc="2.0",
+            id=request_id,
+            error=JsonRpcError(
+                code=error_code,
+                message=error_message,
+                data=error_data
+            )
+        )
+        
+        logger.info(
+            "Translated HTTP error to stdio: status=%s, code=%s, stdio_id=%s",
+            response.status, error_code, request_id
+        )
+        
+        return stdio_message
+    
+    async def translate_sse_event_to_stdio(
+        self,
+        event_data: str,
+        event_type: Optional[str] = None
+    ) -> Optional[JsonRpcMessage]:
+        """Translate Server-Sent Event to stdio JSON-RPC message.
+        
+        SSE events from the HTTP MCP server are translated to JSON-RPC
+        notifications or responses depending on the event type.
+        
+        Args:
+            event_data: SSE event data (JSON string)
+            event_type: Optional SSE event type
+            
+        Returns:
+            JSON-RPC message for stdio, or None if event should be ignored
+            
+        Raises:
+            ProtocolError: If event cannot be translated
+        """
+        logger.debug(
+            "Translating SSE event to stdio: type=%s, data_length=%s",
+            event_type, len(event_data) if event_data else 0
+        )
+        
+        # Parse event data
+        try:
+            event_obj = json.loads(event_data) if event_data else {}
+        except json.JSONDecodeError as e:
+            logger.error("Failed to parse SSE event data: %s", e)
+            raise ProtocolError(
+                "Invalid JSON in SSE event",
+                details={"event_type": event_type, "error": str(e)}
+            )
+        
+        # Check if event is already a JSON-RPC message
+        if "jsonrpc" in event_obj:
+            try:
+                stdio_message = JsonRpcMessage.model_validate(event_obj)
+                logger.info(
+                    "Translated SSE event to stdio: type=%s, method=%s",
+                    event_type, stdio_message.method
+                )
+                return stdio_message
+            except Exception as e:
+                logger.error("Failed to validate JSON-RPC message from SSE: %s", e)
+                raise ProtocolError(
+                    "Invalid JSON-RPC format in SSE event",
+                    details={"event_type": event_type, "error": str(e)}
+                )
+        
+        # Handle different event types
+        if event_type == "message":
+            # Regular message event - wrap as notification
+            stdio_message = JsonRpcMessage(
+                jsonrpc="2.0",
+                method="notifications/message",
+                params=event_obj
+            )
+            logger.info("Translated SSE message event to stdio notification")
+            return stdio_message
+        
+        elif event_type == "error":
+            # Error event - create error notification
+            error_message = event_obj.get("message", "Unknown error")
+            stdio_message = JsonRpcMessage(
+                jsonrpc="2.0",
+                method="notifications/error",
+                params={
+                    "error": error_message,
+                    "details": event_obj
+                }
+            )
+            logger.info("Translated SSE error event to stdio notification")
+            return stdio_message
+        
+        elif event_type in ["ping", "heartbeat"]:
+            # Heartbeat events can be ignored or logged
+            logger.debug("Ignoring SSE heartbeat event")
+            return None
+        
+        else:
+            # Unknown event type - wrap as generic notification
+            stdio_message = JsonRpcMessage(
+                jsonrpc="2.0",
+                method="notifications/event",
+                params={
+                    "type": event_type,
+                    "data": event_obj
+                }
+            )
+            logger.info(
+                "Translated SSE event to stdio notification: type=%s",
+                event_type
+            )
+            return stdio_message
+    
+    def correlate_messages(
+        self,
+        request_id: Union[str, int],
+        response_id: str
+    ) -> None:
+        """Correlate request and response messages.
+        
+        This method is used to manually correlate messages when needed.
+        In most cases, correlation is handled automatically during translation.
+        
+        Args:
+            request_id: Stdio request ID
+            response_id: HTTP response ID
+            
+        Raises:
+            MessageCorrelationError: If correlation fails
+        """
+        # Check if correlation already exists
+        correlation = self._correlation_tracker.get_correlation(request_id)
+        if correlation:
+            logger.warning(
+                "Correlation already exists for request_id=%s, updating",
+                request_id
+            )
+            # Update the HTTP request ID
+            correlation.http_request_id = response_id
+        else:
+            # Create new correlation
+            self._correlation_tracker.add_correlation(
+                stdio_request_id=request_id,
+                http_request_id=response_id,
+                method="unknown"
+            )
+        
+        logger.debug(
+            "Manually correlated messages: stdio_id=%s -> http_id=%s",
+            request_id, response_id
+        )
+    
+    def get_correlation(
+        self,
+        stdio_request_id: Union[str, int]
+    ) -> Optional[MessageCorrelation]:
+        """Get correlation for a stdio request ID.
+        
+        Args:
+            stdio_request_id: Stdio request ID
+            
+        Returns:
+            MessageCorrelation if found, None otherwise
+        """
+        return self._correlation_tracker.get_correlation(stdio_request_id)
+    
+    def mark_correlation_completed(
+        self,
+        stdio_request_id: Union[str, int]
+    ) -> None:
+        """Mark a correlation as completed.
+        
+        Args:
+            stdio_request_id: Stdio request ID
+            
+        Raises:
+            MessageCorrelationError: If correlation not found
+        """
+        self._correlation_tracker.mark_completed(stdio_request_id)
+        logger.debug("Marked correlation completed: stdio_id=%s", stdio_request_id)
+    
+    def mark_correlation_failed(
+        self,
+        stdio_request_id: Union[str, int]
+    ) -> None:
+        """Mark a correlation as failed.
+        
+        Args:
+            stdio_request_id: Stdio request ID
+            
+        Raises:
+            MessageCorrelationError: If correlation not found
+        """
+        self._correlation_tracker.mark_failed(stdio_request_id)
+        logger.debug("Marked correlation failed: stdio_id=%s", stdio_request_id)
+    
+    def clear_correlations(self) -> None:
+        """Clear all correlations.
+        
+        Useful for cleanup and testing.
+        """
+        self._correlation_tracker.clear()
+        logger.debug("Cleared all correlations")
+    
+    def get_pending_correlations(self) -> list[MessageCorrelation]:
+        """Get all pending correlations.
+        
+        Returns:
+            List of pending MessageCorrelation objects
+        """
+        return self._correlation_tracker.get_pending_correlations()