Coverage for src/qdrant_loader_mcp_server/mcp/protocol.py: 95%

1"""MCP Protocol implementation.""" 

2 

3from typing import Any 

4 

5 

6class MCPProtocol: 

7 """MCP Protocol implementation for handling RAG requests.""" 

8 

9 def __init__(self): 

10 """Initialize MCP Protocol.""" 

11 self.version = "2.0" 

12 self.initialized = False 

13 

14 def validate_request(self, request: dict[str, Any]) -> bool: 

15 """Validate MCP request format according to JSON-RPC 2.0 specification. 

16 

17 Args: 

18 request: The request to validate 

19 

20 Returns: 

21 bool: True if request is valid, False otherwise 

22 """ 

23 # Check for required fields 

24 if not isinstance(request, dict): 

25 return False 

26 

27 # Handle empty dict 

28 if not request: 

29 # Allow empty dict only during initialization 

30 return not self.initialized 

31 

32 # For initialization request, be more lenient 

33 if not self.initialized: 

34 if request.get("method") == "initialize": 

35 return True 

36 

37 # Standard validation for other requests 

38 if "jsonrpc" not in request or request["jsonrpc"] != "2.0": 

39 return False 

40 

41 if "method" not in request or not isinstance(request["method"], str): 

42 return False 

43 

44 # For requests (not notifications), id is required 

45 if "id" in request: 

46 if not isinstance(request["id"], str | int): 

47 return False 

48 if request["id"] is None: 

49 return False 

50 else: 

51 # This is a notification, which is valid 

52 return True 

53 

54 # Params is optional but must be object or array if present 

55 if "params" in request and not isinstance(request["params"], dict | list): 

56 return False 

57 

58 return True 

59 

60 def validate_response(self, response: dict[str, Any]) -> bool: 

61 """Validate MCP response format according to JSON-RPC 2.0 specification. 

62 

63 Args: 

64 response: The response to validate 

65 

66 Returns: 

67 bool: True if response is valid, False otherwise 

68 """ 

69 if not isinstance(response, dict): 

70 return False 

71 

72 # Empty response is valid for notifications 

73 if not response: 

74 return True 

75 

76 # Check required fields 

77 if "jsonrpc" not in response or response["jsonrpc"] != "2.0": 

78 return False 

79 

80 if "id" not in response: 

81 return False 

82 

83 if not isinstance(response["id"], str | int): 

84 return False 

85 

86 # Must have either result or error, but not both 

87 has_result = "result" in response 

88 has_error = "error" in response 

89 

90 if not has_result and not has_error: 

91 return False 

92 

93 if has_result and has_error: 

94 return False 

95 

96 # Validate error object structure if present 

97 if has_error: 

98 error = response["error"] 

99 if not isinstance(error, dict): 

100 return False 

101 if "code" not in error or not isinstance(error["code"], int): 

102 return False 

103 if "message" not in error or not isinstance(error["message"], str): 

104 return False 

105 

106 return True 

107 

108 def create_response( 

109 self, 

110 request_id: str | int | None, 

111 result: Any | None = None, 

112 error: dict | None = None, 

113 ) -> dict[str, Any]: 

114 """Create MCP response according to JSON-RPC 2.0 specification. 

115 

116 Args: 

117 request_id: The ID of the request (None for notifications) 

118 result: The result of the request 

119 error: Any error that occurred 

120 

121 Returns: 

122 Dict[str, Any]: The response object 

123 """ 

124 # For notifications, return empty dict 

125 if request_id is None: 

126 return {} 

127 

128 # Create base response 

129 response = {"jsonrpc": self.version, "id": request_id} 

130 

131 # Add either result or error, but not both 

132 if error is not None: 

133 if ( 

134 not isinstance(error, dict) 

135 or "code" not in error 

136 or "message" not in error 

137 ): 

138 error = { 

139 "code": -32603, 

140 "message": "Internal error", 

141 "data": "Invalid error object format", 

142 } 

143 response["error"] = error 

144 else: 

145 # For successful responses, always include result (can be None) 

146 response["result"] = result 

147 

148 # Validate response before returning 

149 if not self.validate_response(response): 

150 return { 

151 "jsonrpc": self.version, 

152 "id": request_id, 

153 "error": { 

154 "code": -32603, 

155 "message": "Internal error", 

156 "data": "Generated invalid response format", 

157 }, 

158 } 

159 

160 return response 

161 

162 def mark_initialized(self): 

163 """Mark the protocol as initialized.""" 

164 self.initialized = True