Coverage for mcpgateway / plugins / framework / hooks / tools.py: 96%

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

2"""Location: ./mcpgateway/plugins/framework/hooks/tools.py 

3Copyright 2025 

4SPDX-License-Identifier: Apache-2.0 

5Authors: Teryl Taylor 

6 

7Pydantic models for tool hooks. 

8""" 

9 

10# Standard 

11from enum import Enum 

12from typing import Any, Optional 

13 

14# Third-Party 

15from pydantic import Field 

16 

17# First-Party 

18from mcpgateway.plugins.framework.hooks.http import HttpHeaderPayload 

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

20 

21 

22class ToolHookType(str, Enum): 

23 """MCP Forge Gateway hook points. 

24 

25 Attributes: 

26 tool_pre_invoke: The tool pre invoke hook. 

27 tool_post_invoke: The tool post invoke hook. 

28 

29 Examples: 

30 >>> ToolHookType.TOOL_PRE_INVOKE 

31 <ToolHookType.TOOL_PRE_INVOKE: 'tool_pre_invoke'> 

32 >>> ToolHookType.TOOL_PRE_INVOKE.value 

33 'tool_pre_invoke' 

34 >>> ToolHookType('tool_post_invoke') 

35 <ToolHookType.TOOL_POST_INVOKE: 'tool_post_invoke'> 

36 >>> list(ToolHookType) 

37 [<ToolHookType.TOOL_PRE_INVOKE: 'tool_pre_invoke'>, <ToolHookType.TOOL_POST_INVOKE: 'tool_post_invoke'>] 

38 """ 

39 

40 TOOL_PRE_INVOKE = "tool_pre_invoke" 

41 TOOL_POST_INVOKE = "tool_post_invoke" 

42 

43 

44class ToolPreInvokePayload(PluginPayload): 

45 """A tool payload for a tool pre-invoke hook. 

46 

47 Args: 

48 name: The tool name. 

49 args: The tool arguments for invocation. 

50 headers: The http pass through headers. 

51 

52 Examples: 

53 >>> payload = ToolPreInvokePayload(name="test_tool", args={"input": "data"}) 

54 >>> payload.name 

55 'test_tool' 

56 >>> payload.args 

57 {'input': 'data'} 

58 >>> payload2 = ToolPreInvokePayload(name="empty") 

59 >>> payload2.args 

60 {} 

61 >>> p = ToolPreInvokePayload(name="calculator", args={"operation": "add", "a": 5, "b": 3}) 

62 >>> p.name 

63 'calculator' 

64 >>> p.args["operation"] 

65 'add' 

66 

67 """ 

68 

69 name: str 

70 args: Optional[dict[str, Any]] = Field(default_factory=dict) 

71 headers: Optional[HttpHeaderPayload] = None 

72 

73 

74class ToolPostInvokePayload(PluginPayload): 

75 """A tool payload for a tool post-invoke hook. 

76 

77 Args: 

78 name: The tool name. 

79 result: The tool invocation result. 

80 

81 Examples: 

82 >>> payload = ToolPostInvokePayload(name="calculator", result={"result": 8, "status": "success"}) 

83 >>> payload.name 

84 'calculator' 

85 >>> payload.result 

86 {'result': 8, 'status': 'success'} 

87 >>> p = ToolPostInvokePayload(name="analyzer", result={"confidence": 0.95, "sentiment": "positive"}) 

88 >>> p.name 

89 'analyzer' 

90 >>> p.result["confidence"] 

91 0.95 

92 """ 

93 

94 name: str 

95 result: Any 

96 

97 

98ToolPreInvokeResult = PluginResult[ToolPreInvokePayload] 

99ToolPostInvokeResult = PluginResult[ToolPostInvokePayload] 

100 

101 

102def _register_tool_hooks() -> None: 

103 """Register Tool hooks in the global registry. 

104 

105 This is called lazily to avoid circular import issues. 

106 """ 

107 # Import here to avoid circular dependency at module load time 

108 # First-Party 

109 from mcpgateway.plugins.framework.hooks.registry import get_hook_registry # pylint: disable=import-outside-toplevel 

110 

111 registry = get_hook_registry() 

112 

113 # Only register if not already registered (idempotent) 

114 if not registry.is_registered(ToolHookType.TOOL_PRE_INVOKE): 114 ↛ exitline 114 didn't return from function '_register_tool_hooks' because the condition on line 114 was always true

115 registry.register_hook(ToolHookType.TOOL_PRE_INVOKE, ToolPreInvokePayload, ToolPreInvokeResult) 

116 registry.register_hook(ToolHookType.TOOL_POST_INVOKE, ToolPostInvokePayload, ToolPostInvokeResult) 

117 

118 

119_register_tool_hooks()