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

mcp_proxy/stdio/jsonrpc.py

@@ -0,0 +1,371 @@
+"""JSON-RPC message parser and serializer for stdio communication.
+
+This module provides functionality for parsing and serializing JSON-RPC 2.0 messages
+with newline-delimited format for stdio communication, along with message correlation
+and ID generation.
+"""
+
+import json
+import uuid
+from typing import Optional, Union, Dict, Any
+from datetime import datetime
+
+from ..models import JsonRpcMessage, JsonRpcError, MessageCorrelation, MessageStatus
+from ..exceptions import InvalidJsonRpcError, MessageCorrelationError
+
+
+class JsonRpcParser:
+    """Parser for JSON-RPC 2.0 messages with newline-delimited format."""
+
+    @staticmethod
+    def parse(line: str) -> JsonRpcMessage:
+        """Parse a newline-delimited JSON-RPC message.
+        
+        Args:
+            line: A single line containing a JSON-RPC message
+            
+        Returns:
+            Parsed JsonRpcMessage object
+            
+        Raises:
+            InvalidJsonRpcError: If the message is invalid or malformed
+        """
+        if not line or not line.strip():
+            raise InvalidJsonRpcError("Empty message line")
+        
+        try:
+            # Remove any trailing newline characters
+            line = line.rstrip('\n\r')
+            
+            # Parse JSON
+            data = json.loads(line)
+            
+            # Validate and create JsonRpcMessage
+            return JsonRpcMessage.model_validate(data)
+            
+        except json.JSONDecodeError as e:
+            raise InvalidJsonRpcError(
+                f"Invalid JSON in message: {e}",
+                details={"line": line, "error": str(e)}
+            )
+        except Exception as e:
+            raise InvalidJsonRpcError(
+                f"Failed to parse JSON-RPC message: {e}",
+                details={"line": line, "error": str(e)}
+            )
+
+    @staticmethod
+    def parse_multiple(lines: str) -> list[JsonRpcMessage]:
+        """Parse multiple newline-delimited JSON-RPC messages.
+        
+        Args:
+            lines: Multiple lines containing JSON-RPC messages
+            
+        Returns:
+            List of parsed JsonRpcMessage objects
+            
+        Raises:
+            InvalidJsonRpcError: If any message is invalid
+        """
+        messages = []
+        for line in lines.split('\n'):
+            if line.strip():  # Skip empty lines
+                messages.append(JsonRpcParser.parse(line))
+        return messages
+
+
+class JsonRpcSerializer:
+    """Serializer for JSON-RPC 2.0 messages with newline-delimited format."""
+
+    @staticmethod
+    def serialize(message: JsonRpcMessage) -> str:
+        """Serialize a JSON-RPC message to newline-delimited format.
+        
+        Args:
+            message: JsonRpcMessage to serialize
+            
+        Returns:
+            Newline-delimited JSON string
+        """
+        # Use Pydantic's JSON serialization with exclude_none
+        json_str = message.model_dump_json(exclude_none=True)
+        
+        # Add newline delimiter
+        return json_str + '\n'
+
+    @staticmethod
+    def serialize_multiple(messages: list[JsonRpcMessage]) -> str:
+        """Serialize multiple JSON-RPC messages to newline-delimited format.
+        
+        Args:
+            messages: List of JsonRpcMessage objects to serialize
+            
+        Returns:
+            Newline-delimited JSON string with multiple messages
+        """
+        return ''.join(JsonRpcSerializer.serialize(msg) for msg in messages)
+
+
+class MessageIdGenerator:
+    """Generator for unique message IDs."""
+
+    def __init__(self, prefix: str = "mcp"):
+        """Initialize the ID generator.
+        
+        Args:
+            prefix: Prefix for generated IDs (default: "mcp")
+        """
+        self.prefix = prefix
+        self._counter = 0
+
+    def generate(self) -> str:
+        """Generate a unique message ID.
+        
+        Returns:
+            Unique message ID string
+        """
+        self._counter += 1
+        # Use UUID for uniqueness combined with counter for ordering
+        return f"{self.prefix}-{self._counter}-{uuid.uuid4().hex[:8]}"
+
+    def reset(self) -> None:
+        """Reset the counter (useful for testing)."""
+        self._counter = 0
+
+
+class MessageCorrelationTracker:
+    """Tracks correlation between stdio and HTTP messages."""
+
+    def __init__(self):
+        """Initialize the correlation tracker."""
+        self._correlations: Dict[Union[str, int], MessageCorrelation] = {}
+
+    def add_correlation(
+        self,
+        stdio_request_id: Union[str, int],
+        http_request_id: str,
+        method: str
+    ) -> MessageCorrelation:
+        """Add a new message correlation.
+        
+        Args:
+            stdio_request_id: ID from the stdio JSON-RPC request
+            http_request_id: ID for the HTTP request
+            method: Method name being called
+            
+        Returns:
+            Created MessageCorrelation object
+            
+        Raises:
+            MessageCorrelationError: If correlation already exists
+        """
+        if stdio_request_id in self._correlations:
+            raise MessageCorrelationError(
+                f"Correlation already exists for stdio request ID: {stdio_request_id}",
+                details={"stdio_request_id": stdio_request_id}
+            )
+
+        correlation = MessageCorrelation(
+            stdio_request_id=stdio_request_id,
+            http_request_id=http_request_id,
+            timestamp=datetime.now(),
+            method=method,
+            status=MessageStatus.PENDING
+        )
+        
+        self._correlations[stdio_request_id] = correlation
+        return correlation
+
+    def get_correlation(
+        self,
+        stdio_request_id: Union[str, int]
+    ) -> Optional[MessageCorrelation]:
+        """Get correlation by stdio request ID.
+        
+        Args:
+            stdio_request_id: ID from the stdio JSON-RPC request
+            
+        Returns:
+            MessageCorrelation object if found, None otherwise
+        """
+        return self._correlations.get(stdio_request_id)
+
+    def get_correlation_by_http_id(
+        self,
+        http_request_id: str
+    ) -> Optional[MessageCorrelation]:
+        """Get correlation by HTTP request ID.
+        
+        Args:
+            http_request_id: ID from the HTTP request
+            
+        Returns:
+            MessageCorrelation object if found, None otherwise
+        """
+        for correlation in self._correlations.values():
+            if correlation.http_request_id == http_request_id:
+                return correlation
+        return None
+
+    def mark_completed(self, stdio_request_id: Union[str, int]) -> None:
+        """Mark a correlation as completed.
+        
+        Args:
+            stdio_request_id: ID from the stdio JSON-RPC request
+            
+        Raises:
+            MessageCorrelationError: If correlation not found
+        """
+        correlation = self.get_correlation(stdio_request_id)
+        if not correlation:
+            raise MessageCorrelationError(
+                f"No correlation found for stdio request ID: {stdio_request_id}",
+                details={"stdio_request_id": stdio_request_id}
+            )
+        correlation.mark_completed()
+
+    def mark_failed(self, stdio_request_id: Union[str, int]) -> None:
+        """Mark a correlation as failed.
+        
+        Args:
+            stdio_request_id: ID from the stdio JSON-RPC request
+            
+        Raises:
+            MessageCorrelationError: If correlation not found
+        """
+        correlation = self.get_correlation(stdio_request_id)
+        if not correlation:
+            raise MessageCorrelationError(
+                f"No correlation found for stdio request ID: {stdio_request_id}",
+                details={"stdio_request_id": stdio_request_id}
+            )
+        correlation.mark_failed()
+
+    def remove_correlation(self, stdio_request_id: Union[str, int]) -> None:
+        """Remove a correlation from tracking.
+        
+        Args:
+            stdio_request_id: ID from the stdio JSON-RPC request
+        """
+        self._correlations.pop(stdio_request_id, None)
+
+    def get_pending_correlations(self) -> list[MessageCorrelation]:
+        """Get all pending correlations.
+        
+        Returns:
+            List of pending MessageCorrelation objects
+        """
+        return [
+            corr for corr in self._correlations.values()
+            if corr.is_pending()
+        ]
+
+    def clear(self) -> None:
+        """Clear all correlations (useful for testing and cleanup)."""
+        self._correlations.clear()
+
+
+def create_request(
+    method: str,
+    params: Optional[Any] = None,
+    request_id: Optional[Union[str, int]] = None
+) -> JsonRpcMessage:
+    """Create a JSON-RPC request message.
+    
+    Args:
+        method: Method name to call
+        params: Method parameters (optional)
+        request_id: Request ID (optional, will be generated if not provided)
+        
+    Returns:
+        JsonRpcMessage request object
+    """
+    if request_id is None:
+        # Generate a unique ID
+        request_id = MessageIdGenerator().generate()
+    
+    return JsonRpcMessage(
+        jsonrpc="2.0",
+        id=request_id,
+        method=method,
+        params=params
+    )
+
+
+def create_response(
+    request_id: Union[str, int],
+    result: Optional[Any] = None,
+    error: Optional[JsonRpcError] = None
+) -> JsonRpcMessage:
+    """Create a JSON-RPC response message.
+    
+    Args:
+        request_id: ID from the original request
+        result: Result data (for success response)
+        error: Error object (for error response)
+        
+    Returns:
+        JsonRpcMessage response object
+        
+    Raises:
+        ValueError: If both result and error are provided, or neither
+    """
+    if (result is None and error is None) or (result is not None and error is not None):
+        raise ValueError("Must provide either result or error, but not both")
+    
+    return JsonRpcMessage(
+        jsonrpc="2.0",
+        id=request_id,
+        result=result,
+        error=error
+    )
+
+
+def create_notification(
+    method: str,
+    params: Optional[Any] = None
+) -> JsonRpcMessage:
+    """Create a JSON-RPC notification message (request without ID).
+    
+    Args:
+        method: Method name to call
+        params: Method parameters (optional)
+        
+    Returns:
+        JsonRpcMessage notification object
+    """
+    return JsonRpcMessage(
+        jsonrpc="2.0",
+        method=method,
+        params=params
+    )
+
+
+def create_error_response(
+    request_id: Union[str, int],
+    code: int,
+    message: str,
+    data: Optional[Any] = None
+) -> JsonRpcMessage:
+    """Create a JSON-RPC error response message.
+    
+    Args:
+        request_id: ID from the original request
+        code: Error code
+        message: Error message
+        data: Additional error data (optional)
+        
+    Returns:
+        JsonRpcMessage error response object
+    """
+    error = JsonRpcError(
+        code=code,
+        message=message,
+        data=data
+    )
+    
+    return JsonRpcMessage(
+        jsonrpc="2.0",
+        id=request_id,
+        error=error
+    )

mcp_proxy/translator/__init__.py

@@ -0,0 +1,11 @@
+"""Protocol translator module for MCP Proxy.
+
+This module handles translation between stdio and HTTP MCP protocols.
+"""
+
+try:
+    from .translator import ProtocolTranslatorImpl
+    __all__ = ["ProtocolTranslatorImpl"]
+except ImportError:
+    # Module not yet implemented
+    __all__ = []