Coverage for mcpgateway / plugins / framework / decorator.py: 100%

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

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

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Hook decorator for dynamically registering plugin hooks. 

8 

9This module provides decorators for marking plugin methods as hook handlers. 

10Plugins can use these decorators to: 

111. Override the default hook naming convention 

122. Register custom hooks not in the standard framework 

13 

14Examples: 

15 Override hook method name:: 

16 

17 class MyPlugin(Plugin): 

18 @hook(ToolHookType.TOOL_PRE_INVOKE) 

19 def custom_name_for_tool_hook(self, payload, context): 

20 # This gets called for tool_pre_invoke even though 

21 # the method name doesn't match 

22 return ToolPreInvokeResult(continue_processing=True) 

23 

24 Register a completely new hook type:: 

25 

26 class MyPlugin(Plugin): 

27 @hook("custom_pre_process", CustomPayload, CustomResult) 

28 def my_custom_hook(self, payload, context): 

29 # This registers a new hook type dynamically 

30 return CustomResult(continue_processing=True) 

31 

32 Use default convention (no decorator needed):: 

33 

34 class MyPlugin(Plugin): 

35 def tool_pre_invoke(self, payload, context): 

36 # Automatically recognized by naming convention 

37 return ToolPreInvokeResult(continue_processing=True) 

38""" 

39 

40# Standard 

41from typing import Callable, Optional, Type, TypeVar 

42 

43# Third-Party 

44from pydantic import BaseModel 

45 

46# First-Party 

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

48 

49# Attribute name for storing hook metadata on functions 

50_HOOK_METADATA_ATTR = "_plugin_hook_metadata" 

51 

52# Type vars for type hints 

53P = TypeVar("P", bound=PluginPayload) # Payload type 

54R = TypeVar("R", bound=PluginResult) # Result type 

55 

56 

57class HookMetadata: 

58 """Metadata stored on decorated hook methods. 

59 

60 Attributes: 

61 hook_type: The hook type identifier (e.g., 'tool_pre_invoke') 

62 payload_type: Optional payload class for hook registration 

63 result_type: Optional result class for hook registration 

64 """ 

65 

66 def __init__( 

67 self, 

68 hook_type: str, 

69 payload_type: Optional[Type[BaseModel]] = None, 

70 result_type: Optional[Type[BaseModel]] = None, 

71 ): 

72 """Initialize hook metadata. 

73 

74 Args: 

75 hook_type: The hook type identifier 

76 payload_type: Optional payload class for registering new hooks 

77 result_type: Optional result class for registering new hooks 

78 """ 

79 self.hook_type = hook_type 

80 self.payload_type = payload_type 

81 self.result_type = result_type 

82 

83 

84def hook( 

85 hook_type: str, 

86 payload_type: Optional[Type[P]] = None, 

87 result_type: Optional[Type[R]] = None, 

88) -> Callable[[Callable], Callable]: 

89 """Decorator to mark a method as a plugin hook handler. 

90 

91 This decorator attaches metadata to a method so the Plugin class can 

92 discover it during initialization and register it with the appropriate 

93 hook type. 

94 

95 Args: 

96 hook_type: The hook type identifier (e.g., 'tool_pre_invoke') 

97 payload_type: Optional payload class for registering new hook types 

98 result_type: Optional result class for registering new hook types 

99 

100 Returns: 

101 Decorator function that marks the method with hook metadata 

102 

103 Examples: 

104 Override method name:: 

105 

106 @hook(ToolHookType.TOOL_PRE_INVOKE) 

107 def my_custom_method_name(self, payload, context): 

108 return ToolPreInvokeResult(continue_processing=True) 

109 

110 Register new hook type:: 

111 

112 @hook("email_pre_send", EmailPayload, EmailResult) 

113 def handle_email(self, payload, context): 

114 return EmailResult(continue_processing=True) 

115 """ 

116 

117 def decorator(func: Callable) -> Callable: 

118 """Inner decorator that attaches metadata to the function. 

119 

120 Args: 

121 func: The function to decorate 

122 

123 Returns: 

124 The same function with metadata attached 

125 """ 

126 # Store metadata on the function object 

127 metadata = HookMetadata(hook_type, payload_type, result_type) 

128 setattr(func, _HOOK_METADATA_ATTR, metadata) 

129 return func 

130 

131 return decorator 

132 

133 

134def get_hook_metadata(func: Callable) -> Optional[HookMetadata]: 

135 """Get hook metadata from a decorated function. 

136 

137 Args: 

138 func: The function to check 

139 

140 Returns: 

141 HookMetadata if the function is decorated, None otherwise 

142 

143 Examples: 

144 >>> @hook("test_hook") 

145 ... def test_func(): 

146 ... pass 

147 >>> metadata = get_hook_metadata(test_func) 

148 >>> metadata.hook_type 

149 'test_hook' 

150 >>> get_hook_metadata(lambda: None) is None 

151 True 

152 """ 

153 return getattr(func, _HOOK_METADATA_ATTR, None) 

154 

155 

156def has_hook_metadata(func: Callable) -> bool: 

157 """Check if a function has hook metadata. 

158 

159 Args: 

160 func: The function to check 

161 

162 Returns: 

163 True if the function is decorated with @hook, False otherwise 

164 

165 Examples: 

166 >>> @hook("test_hook") 

167 ... def decorated(): 

168 ... pass 

169 >>> has_hook_metadata(decorated) 

170 True 

171 >>> has_hook_metadata(lambda: None) 

172 False 

173 """ 

174 return hasattr(func, _HOOK_METADATA_ATTR)