Coverage for mcpgateway / plugins / framework / external / proto_convert.py: 100%

1# -*- coding: utf-8 -*- 

2"""Location: ./mcpgateway/plugins/framework/external/proto_convert.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Conversion utilities between Pydantic models and protobuf messages. 

8 

9This module provides efficient conversion functions that use explicit protobuf 

10messages where possible, falling back to Struct for dynamic fields. 

11""" 

12 

13# pylint: disable=no-member 

14 

15# Standard 

16 

17# Third-Party 

18from google.protobuf import json_format 

19 

20# First-Party 

21from mcpgateway.plugins.framework.external.grpc.proto import plugin_service_pb2 

22from mcpgateway.plugins.framework.models import ( 

23 PluginResult, 

24) 

25from mcpgateway.plugins.framework.models import GlobalContext as PydanticGlobalContext 

26from mcpgateway.plugins.framework.models import PluginContext as PydanticPluginContext 

27from mcpgateway.plugins.framework.models import PluginViolation as PydanticPluginViolation 

28 

29 

30def pydantic_global_context_to_proto(ctx: PydanticGlobalContext) -> plugin_service_pb2.GlobalContext: 

31 """Convert Pydantic GlobalContext to protobuf GlobalContext. 

32 

33 Args: 

34 ctx: The Pydantic GlobalContext model. 

35 

36 Returns: 

37 The protobuf GlobalContext message. 

38 """ 

39 proto_ctx = plugin_service_pb2.GlobalContext( 

40 request_id=ctx.request_id, 

41 server_id=ctx.server_id or "", 

42 tenant_id=ctx.tenant_id or "", 

43 ) 

44 

45 # Handle user field (can be string or dict) 

46 if ctx.user is not None: 

47 if isinstance(ctx.user, str): 

48 proto_ctx.user_string = ctx.user 

49 elif isinstance(ctx.user, dict): 

50 json_format.ParseDict(ctx.user, proto_ctx.user_struct) 

51 

52 # Handle dynamic fields with Struct 

53 if ctx.metadata: 

54 json_format.ParseDict(ctx.metadata, proto_ctx.metadata) 

55 if ctx.state: 

56 json_format.ParseDict(ctx.state, proto_ctx.state) 

57 

58 return proto_ctx 

59 

60 

61def proto_global_context_to_pydantic(proto_ctx: plugin_service_pb2.GlobalContext) -> PydanticGlobalContext: 

62 """Convert protobuf GlobalContext to Pydantic GlobalContext. 

63 

64 Args: 

65 proto_ctx: The protobuf GlobalContext message. 

66 

67 Returns: 

68 The Pydantic GlobalContext model. 

69 """ 

70 # Handle user field 

71 user = None 

72 if proto_ctx.HasField("user_string"): 

73 user = proto_ctx.user_string 

74 elif proto_ctx.HasField("user_struct"): 

75 user = json_format.MessageToDict(proto_ctx.user_struct) 

76 

77 return PydanticGlobalContext( 

78 request_id=proto_ctx.request_id, 

79 server_id=proto_ctx.server_id or None, 

80 tenant_id=proto_ctx.tenant_id or None, 

81 user=user, 

82 metadata=json_format.MessageToDict(proto_ctx.metadata) if proto_ctx.metadata.fields else {}, 

83 state=json_format.MessageToDict(proto_ctx.state) if proto_ctx.state.fields else {}, 

84 ) 

85 

86 

87def pydantic_context_to_proto(ctx: PydanticPluginContext) -> plugin_service_pb2.PluginContext: 

88 """Convert Pydantic PluginContext to protobuf PluginContext. 

89 

90 Args: 

91 ctx: The Pydantic PluginContext model. 

92 

93 Returns: 

94 The protobuf PluginContext message. 

95 """ 

96 proto_ctx = plugin_service_pb2.PluginContext( 

97 global_context=pydantic_global_context_to_proto(ctx.global_context), 

98 ) 

99 

100 if ctx.state: 

101 json_format.ParseDict(ctx.state, proto_ctx.state) 

102 if ctx.metadata: 

103 json_format.ParseDict(ctx.metadata, proto_ctx.metadata) 

104 

105 return proto_ctx 

106 

107 

108def proto_context_to_pydantic(proto_ctx: plugin_service_pb2.PluginContext) -> PydanticPluginContext: 

109 """Convert protobuf PluginContext to Pydantic PluginContext. 

110 

111 Args: 

112 proto_ctx: The protobuf PluginContext message. 

113 

114 Returns: 

115 The Pydantic PluginContext model. 

116 """ 

117 return PydanticPluginContext( 

118 global_context=proto_global_context_to_pydantic(proto_ctx.global_context), 

119 state=json_format.MessageToDict(proto_ctx.state) if proto_ctx.state.fields else {}, 

120 metadata=json_format.MessageToDict(proto_ctx.metadata) if proto_ctx.metadata.fields else {}, 

121 ) 

122 

123 

124def proto_context_to_dict(proto_ctx: plugin_service_pb2.PluginContext) -> dict: 

125 """Convert protobuf PluginContext directly to dict (for server use). 

126 

127 This avoids the intermediate Pydantic model when only a dict is needed. 

128 

129 Args: 

130 proto_ctx: The protobuf PluginContext message. 

131 

132 Returns: 

133 Dictionary representation of the context. 

134 """ 

135 gc = proto_ctx.global_context 

136 

137 # Handle user field 

138 user = None 

139 if gc.HasField("user_string"): 

140 user = gc.user_string 

141 elif gc.HasField("user_struct"): 

142 user = json_format.MessageToDict(gc.user_struct) 

143 

144 return { 

145 "global_context": { 

146 "request_id": gc.request_id, 

147 "server_id": gc.server_id or None, 

148 "tenant_id": gc.tenant_id or None, 

149 "user": user, 

150 "metadata": json_format.MessageToDict(gc.metadata) if gc.metadata.fields else {}, 

151 "state": json_format.MessageToDict(gc.state) if gc.state.fields else {}, 

152 }, 

153 "state": json_format.MessageToDict(proto_ctx.state) if proto_ctx.state.fields else {}, 

154 "metadata": json_format.MessageToDict(proto_ctx.metadata) if proto_ctx.metadata.fields else {}, 

155 } 

156 

157 

158def pydantic_violation_to_proto(violation: PydanticPluginViolation) -> plugin_service_pb2.PluginViolation: 

159 """Convert Pydantic PluginViolation to protobuf PluginViolation. 

160 

161 Args: 

162 violation: The Pydantic PluginViolation model. 

163 

164 Returns: 

165 The protobuf PluginViolation message. 

166 """ 

167 proto_violation = plugin_service_pb2.PluginViolation( 

168 reason=violation.reason, 

169 description=violation.description, 

170 code=violation.code, 

171 plugin_name=violation.plugin_name or "", 

172 mcp_error_code=violation.mcp_error_code or 0, 

173 ) 

174 

175 if violation.details: 

176 json_format.ParseDict(violation.details, proto_violation.details) 

177 

178 return proto_violation 

179 

180 

181def proto_violation_to_pydantic(proto_violation: plugin_service_pb2.PluginViolation) -> PydanticPluginViolation: 

182 """Convert protobuf PluginViolation to Pydantic PluginViolation. 

183 

184 Args: 

185 proto_violation: The protobuf PluginViolation message. 

186 

187 Returns: 

188 The Pydantic PluginViolation model. 

189 """ 

190 violation = PydanticPluginViolation( 

191 reason=proto_violation.reason, 

192 description=proto_violation.description, 

193 code=proto_violation.code, 

194 details=json_format.MessageToDict(proto_violation.details) if proto_violation.details.fields else {}, 

195 mcp_error_code=proto_violation.mcp_error_code if proto_violation.mcp_error_code else None, 

196 ) 

197 if proto_violation.plugin_name: 

198 violation.plugin_name = proto_violation.plugin_name 

199 return violation 

200 

201 

202def pydantic_result_to_proto_base(result: PluginResult) -> plugin_service_pb2.PluginResultBase: 

203 """Convert common PluginResult fields to protobuf PluginResultBase. 

204 

205 Args: 

206 result: The Pydantic PluginResult model. 

207 

208 Returns: 

209 The protobuf PluginResultBase message with common fields. 

210 """ 

211 proto_result = plugin_service_pb2.PluginResultBase( 

212 continue_processing=result.continue_processing, 

213 ) 

214 

215 if result.violation: 

216 proto_result.violation.CopyFrom(pydantic_violation_to_proto(result.violation)) 

217 

218 if result.metadata: 

219 json_format.ParseDict(result.metadata, proto_result.metadata) 

220 

221 return proto_result 

222 

223 

224def update_pydantic_result_from_proto_base( 

225 result: PluginResult, 

226 proto_base: plugin_service_pb2.PluginResultBase, 

227) -> None: 

228 """Update a Pydantic PluginResult with values from PluginResultBase. 

229 

230 Args: 

231 result: The Pydantic PluginResult to update. 

232 proto_base: The protobuf PluginResultBase with common fields. 

233 """ 

234 result.continue_processing = proto_base.continue_processing 

235 

236 if proto_base.HasField("violation"): 

237 result.violation = proto_violation_to_pydantic(proto_base.violation) 

238 

239 if proto_base.metadata.fields: 

240 result.metadata = json_format.MessageToDict(proto_base.metadata) 

241 

242 

243def update_pydantic_context_from_proto( 

244 ctx: PydanticPluginContext, 

245 proto_ctx: plugin_service_pb2.PluginContext, 

246) -> None: 

247 """Update a Pydantic PluginContext in-place from protobuf PluginContext. 

248 

249 Args: 

250 ctx: The Pydantic PluginContext to update. 

251 proto_ctx: The protobuf PluginContext with updated values. 

252 """ 

253 ctx.state = json_format.MessageToDict(proto_ctx.state) if proto_ctx.state.fields else {} 

254 ctx.metadata = json_format.MessageToDict(proto_ctx.metadata) if proto_ctx.metadata.fields else {} 

255 

256 # Update global context state 

257 if proto_ctx.global_context.state.fields: 

258 ctx.global_context.state = json_format.MessageToDict(proto_ctx.global_context.state)