Coverage for mcpgateway / plugins / framework / hooks / registry.py: 100%

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

2"""Location: ./mcpgateway/plugins/framework/hook_registry.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Hook Registry. 

8This module provides a global registry for mapping hook types to their 

9corresponding payload and result Pydantic models. This enables external 

10plugins to properly serialize/deserialize payloads without needing direct 

11access to the specific plugin implementations. 

12""" 

13 

14# Standard 

15from typing import Dict, Optional, Type, Union 

16 

17# First-Party 

18from mcpgateway.plugins.framework.models import PluginPayload, PluginResult 

19 

20 

21class HookRegistry: 

22 """Global registry for hook type metadata. 

23 

24 This singleton registry maintains mappings between hook type names and their 

25 associated Pydantic models for payloads and results. It enables dynamic 

26 serialization/deserialization for external plugins. 

27 

28 Examples: 

29 >>> from mcpgateway.plugins.framework import PluginPayload, PluginResult 

30 >>> registry = HookRegistry() 

31 >>> registry.register_hook("test_hook", PluginPayload, PluginResult) 

32 >>> registry.get_payload_type("test_hook") 

33 <class 'pydantic.main.BaseModel'> 

34 >>> registry.get_result_type("test_hook") 

35 <class 'mcpgateway.plugins.framework.models.PluginResult'> 

36 """ 

37 

38 _instance: Optional["HookRegistry"] = None 

39 _hook_payloads: Dict[str, Type[PluginPayload]] = {} 

40 _hook_results: Dict[str, Type[PluginResult]] = {} 

41 

42 def __new__(cls) -> "HookRegistry": 

43 """Ensure singleton pattern for the registry. 

44 

45 Returns: 

46 The singleton HookRegistry instance. 

47 """ 

48 if cls._instance is None: 

49 cls._instance = super().__new__(cls) 

50 return cls._instance 

51 

52 def register_hook( 

53 self, 

54 hook_type: str, 

55 payload_class: Type[PluginPayload], 

56 result_class: Type[PluginResult], 

57 ) -> None: 

58 """Register a hook type with its payload and result classes. 

59 

60 Args: 

61 hook_type: The hook type identifier (e.g., "prompt_pre_fetch"). 

62 payload_class: The Pydantic model class for the hook's payload. 

63 result_class: The Pydantic model class for the hook's result. 

64 

65 Examples: 

66 >>> registry = HookRegistry() 

67 >>> from mcpgateway.plugins.framework import PluginPayload, PluginResult 

68 >>> registry.register_hook("custom_hook", PluginPayload, PluginResult) 

69 """ 

70 self._hook_payloads[hook_type] = payload_class 

71 self._hook_results[hook_type] = result_class 

72 

73 def get_payload_type(self, hook_type: str) -> Optional[Type[PluginPayload]]: 

74 """Get the payload class for a hook type. 

75 

76 Args: 

77 hook_type: The hook type identifier. 

78 

79 Returns: 

80 The Pydantic payload class, or None if not registered. 

81 

82 Examples: 

83 >>> registry = HookRegistry() 

84 >>> registry.get_payload_type("unknown_hook") 

85 """ 

86 return self._hook_payloads.get(hook_type) 

87 

88 def get_result_type(self, hook_type: str) -> Optional[Type[PluginResult]]: 

89 """Get the result class for a hook type. 

90 

91 Args: 

92 hook_type: The hook type identifier. 

93 

94 Returns: 

95 The Pydantic result class, or None if not registered. 

96 

97 Examples: 

98 >>> registry = HookRegistry() 

99 >>> registry.get_result_type("unknown_hook") 

100 """ 

101 return self._hook_results.get(hook_type) 

102 

103 def json_to_payload(self, hook_type: str, payload: Union[str, dict]) -> PluginPayload: 

104 """Convert JSON to the appropriate payload Pydantic model. 

105 

106 Args: 

107 hook_type: The hook type identifier. 

108 payload: The payload as JSON string or dictionary. 

109 

110 Returns: 

111 The deserialized Pydantic payload object. 

112 

113 Raises: 

114 ValueError: If the hook type is not registered. 

115 

116 Examples: 

117 >>> registry = HookRegistry() 

118 >>> from mcpgateway.plugins.framework.hooks.prompts import PromptPrehookPayload, PromptPrehookResult 

119 >>> registry.register_hook("test", PromptPrehookPayload, PromptPrehookResult) 

120 >>> payload = registry.json_to_payload("test", {"prompt_id": "123"}) 

121 """ 

122 payload_class = self.get_payload_type(hook_type) 

123 if not payload_class: 

124 raise ValueError(f"No payload type registered for hook: {hook_type}") 

125 

126 if isinstance(payload, str): 

127 return payload_class.model_validate_json(payload) 

128 return payload_class.model_validate(payload) 

129 

130 def json_to_result(self, hook_type: str, result: Union[str, dict]) -> PluginResult: 

131 """Convert JSON to the appropriate result Pydantic model. 

132 

133 Args: 

134 hook_type: The hook type identifier. 

135 result: The result as JSON string or dictionary. 

136 

137 Returns: 

138 The deserialized Pydantic result object. 

139 

140 Raises: 

141 ValueError: If the hook type is not registered. 

142 

143 Examples: 

144 >>> registry = HookRegistry() 

145 >>> from mcpgateway.plugins.framework import PluginPayload, PluginResult 

146 >>> registry.register_hook("test", PluginPayload, PluginResult) 

147 >>> result = registry.json_to_result("test", '{"continue_processing": true}') 

148 """ 

149 result_class = self.get_result_type(hook_type) 

150 if not result_class: 

151 raise ValueError(f"No result type registered for hook: {hook_type}") 

152 

153 if isinstance(result, str): 

154 return result_class.model_validate_json(result) 

155 return result_class.model_validate(result) 

156 

157 def is_registered(self, hook_type: str) -> bool: 

158 """Check if a hook type is registered. 

159 

160 Args: 

161 hook_type: The hook type identifier. 

162 

163 Returns: 

164 True if the hook is registered, False otherwise. 

165 

166 Examples: 

167 >>> registry = HookRegistry() 

168 >>> registry.is_registered("unknown") 

169 False 

170 """ 

171 return hook_type in self._hook_payloads and hook_type in self._hook_results 

172 

173 def get_registered_hooks(self) -> list[str]: 

174 """Get all registered hook types. 

175 

176 Returns: 

177 List of registered hook type identifiers. 

178 

179 Examples: 

180 >>> registry = HookRegistry() 

181 >>> hooks = registry.get_registered_hooks() 

182 >>> isinstance(hooks, list) 

183 True 

184 """ 

185 return list(self._hook_payloads.keys()) 

186 

187 

188# Global singleton instance 

189_global_registry = HookRegistry() 

190 

191 

192def get_hook_registry() -> HookRegistry: 

193 """Get the global hook registry instance. 

194 

195 Returns: 

196 The singleton HookRegistry instance. 

197 

198 Examples: 

199 >>> registry = get_hook_registry() 

200 >>> isinstance(registry, HookRegistry) 

201 True 

202 """ 

203 return _global_registry