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

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# pylint: disable=no-member 

13 

14# Standard 

15 

16# Third-Party 

17from google.protobuf import json_format 

18 

19# First-Party 

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

21from mcpgateway.plugins.framework.models import ( 

22 PluginResult, 

23) 

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

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

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

27 

28 

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

30 """Convert Pydantic GlobalContext to protobuf GlobalContext. 

31 

32 Args: 

33 ctx: The Pydantic GlobalContext model. 

34 

35 Returns: 

36 The protobuf GlobalContext message. 

37 """ 

38 proto_ctx = plugin_service_pb2.GlobalContext( 

39 request_id=ctx.request_id, 

40 server_id=ctx.server_id or "", 

41 tenant_id=ctx.tenant_id or "", 

42 ) 

43 

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

45 if ctx.user is not None: 

46 if isinstance(ctx.user, str): 

47 proto_ctx.user_string = ctx.user 

48 elif isinstance(ctx.user, dict): 48 ↛ 52line 48 didn't jump to line 52 because the condition on line 48 was always true

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

50 

51 # Handle dynamic fields with Struct 

52 if ctx.metadata: 

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

54 if ctx.state: 

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

56 

57 return proto_ctx 

58 

59 

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

61 """Convert protobuf GlobalContext to Pydantic GlobalContext. 

62 

63 Args: 

64 proto_ctx: The protobuf GlobalContext message. 

65 

66 Returns: 

67 The Pydantic GlobalContext model. 

68 """ 

69 # Handle user field 

70 user = None 

71 if proto_ctx.HasField("user_string"): 

72 user = proto_ctx.user_string 

73 elif proto_ctx.HasField("user_struct"): 

74 user = json_format.MessageToDict(proto_ctx.user_struct) 

75 

76 return PydanticGlobalContext( 

77 request_id=proto_ctx.request_id, 

78 server_id=proto_ctx.server_id or None, 

79 tenant_id=proto_ctx.tenant_id or None, 

80 user=user, 

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

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

83 ) 

84 

85 

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

87 """Convert Pydantic PluginContext to protobuf PluginContext. 

88 

89 Args: 

90 ctx: The Pydantic PluginContext model. 

91 

92 Returns: 

93 The protobuf PluginContext message. 

94 """ 

95 proto_ctx = plugin_service_pb2.PluginContext( 

96 global_context=pydantic_global_context_to_proto(ctx.global_context), 

97 ) 

98 

99 if ctx.state: 

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

101 if ctx.metadata: 

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

103 

104 return proto_ctx 

105 

106 

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

108 """Convert protobuf PluginContext to Pydantic PluginContext. 

109 

110 Args: 

111 proto_ctx: The protobuf PluginContext message. 

112 

113 Returns: 

114 The Pydantic PluginContext model. 

115 """ 

116 return PydanticPluginContext( 

117 global_context=proto_global_context_to_pydantic(proto_ctx.global_context), 

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

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

120 ) 

121 

122 

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

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

125 

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

127 

128 Args: 

129 proto_ctx: The protobuf PluginContext message. 

130 

131 Returns: 

132 Dictionary representation of the context. 

133 """ 

134 gc = proto_ctx.global_context 

135 

136 # Handle user field 

137 user = None 

138 if gc.HasField("user_string"): 

139 user = gc.user_string 

140 elif gc.HasField("user_struct"): 

141 user = json_format.MessageToDict(gc.user_struct) 

142 

143 return { 

144 "global_context": { 

145 "request_id": gc.request_id, 

146 "server_id": gc.server_id or None, 

147 "tenant_id": gc.tenant_id or None, 

148 "user": user, 

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

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

151 }, 

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

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

154 } 

155 

156 

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

158 """Convert Pydantic PluginViolation to protobuf PluginViolation. 

159 

160 Args: 

161 violation: The Pydantic PluginViolation model. 

162 

163 Returns: 

164 The protobuf PluginViolation message. 

165 """ 

166 proto_violation = plugin_service_pb2.PluginViolation( 

167 reason=violation.reason, 

168 description=violation.description, 

169 code=violation.code, 

170 plugin_name=violation.plugin_name or "", 

171 mcp_error_code=violation.mcp_error_code or 0, 

172 ) 

173 

174 if violation.details: 

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

176 

177 return proto_violation 

178 

179 

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

181 """Convert protobuf PluginViolation to Pydantic PluginViolation. 

182 

183 Args: 

184 proto_violation: The protobuf PluginViolation message. 

185 

186 Returns: 

187 The Pydantic PluginViolation model. 

188 """ 

189 violation = PydanticPluginViolation( 

190 reason=proto_violation.reason, 

191 description=proto_violation.description, 

192 code=proto_violation.code, 

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

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

195 ) 

196 if proto_violation.plugin_name: 

197 violation.plugin_name = proto_violation.plugin_name 

198 return violation 

199 

200 

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

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

203 

204 Args: 

205 result: The Pydantic PluginResult model. 

206 

207 Returns: 

208 The protobuf PluginResultBase message with common fields. 

209 """ 

210 proto_result = plugin_service_pb2.PluginResultBase( 

211 continue_processing=result.continue_processing, 

212 ) 

213 

214 if result.violation: 

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

216 

217 if result.metadata: 

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

219 

220 return proto_result 

221 

222 

223def update_pydantic_result_from_proto_base( 

224 result: PluginResult, 

225 proto_base: plugin_service_pb2.PluginResultBase, 

226) -> None: 

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

228 

229 Args: 

230 result: The Pydantic PluginResult to update. 

231 proto_base: The protobuf PluginResultBase with common fields. 

232 """ 

233 result.continue_processing = proto_base.continue_processing 

234 

235 if proto_base.HasField("violation"): 

236 result.violation = proto_violation_to_pydantic(proto_base.violation) 

237 

238 if proto_base.metadata.fields: 

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

240 

241 

242def update_pydantic_context_from_proto( 

243 ctx: PydanticPluginContext, 

244 proto_ctx: plugin_service_pb2.PluginContext, 

245) -> None: 

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

247 

248 Args: 

249 ctx: The Pydantic PluginContext to update. 

250 proto_ctx: The protobuf PluginContext with updated values. 

251 """ 

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

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

254 

255 # Update global context state 

256 if proto_ctx.global_context.state.fields: 

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