feat: Add tool_pre_invoke and tool_post_invoke plugin hooks (#686)

* feat: add tool hooks to plugin system

Signed-off-by: Shamsul Arefin <[email protected]>

* feat: connect plugins to tool invoke, pii_filter improvements

Signed-off-by: Shamsul Arefin <[email protected]>

* Add  PluginViolationError: If plugin blocks tool invocation. to the docstring's Raises section.

Signed-off-by: Shamsul Arefin <[email protected]>

* Apply linter fixes

* Rebase and lint

Signed-off-by: Mihai Criveti <[email protected]>

* Update logging and docs

Signed-off-by: Mihai Criveti <[email protected]>

---------

Signed-off-by: Shamsul Arefin <[email protected]>
Signed-off-by: Mihai Criveti <[email protected]>
Co-authored-by: Shamsul Arefin <[email protected]>
Co-authored-by: Mihai Criveti <[email protected]>

Author: 3 people

docs/docs/using/plugins/index.md

@@ -94,7 +94,7 @@ are defined as follows:
 | **description** | The description of the plugin configuration. | A plugin for replacing bad words. |
 | **version** | The version of the plugin configuration. | 0.1 |
 | **author** | The team that wrote the plugin. | MCP Context Forge |
-| **hooks** | A list of hooks for which the plugin will be executed. **Note**: currently supports two hooks: "prompt_pre_fetch", "prompt_post_fetch"  | ["prompt_pre_fetch", "prompt_post_fetch"] |
+| **hooks** | A list of hooks for which the plugin will be executed. Supported hooks: "prompt_pre_fetch", "prompt_post_fetch", "tool_pre_invoke", "tool_post_invoke"  | ["prompt_pre_fetch", "prompt_post_fetch", "tool_pre_invoke", "tool_post_invoke"] |
 | **tags** | Descriptive keywords that make the configuration searchable. | ["security", "filter"] |
 | **mode** | Mode of operation of the plugin. - enforce (stops during a violation), permissive (audits a violation but doesn't stop), disabled (disabled) | permissive |
 | **priority** | The priority in which the plugin will run - 0 is higher priority | 100 |
@@ -163,10 +163,25 @@ Currently implemented hooks:
 |------|-------------|-----------|
 | `prompt_pre_fetch` | Before prompt retrieval | Validate/modify prompt arguments |
 | `prompt_post_fetch` | After prompt rendering | Filter/transform rendered prompts |
+| `tool_pre_invoke` | Before tool invocation | Validate/modify tool arguments, block dangerous operations |
+| `tool_post_invoke` | After tool execution | Filter/transform tool results, audit tool usage |
+
+### Tool Hooks Details
+
+The tool hooks enable plugins to intercept and modify tool invocations:
+
+- **`tool_pre_invoke`**: Receives the tool name and arguments before execution. Can modify arguments or block the invocation entirely.
+- **`tool_post_invoke`**: Receives the tool result after execution. Can modify the result or block it from being returned.
+
+Example use cases:
+- PII detection and masking in tool inputs/outputs
+- Rate limiting specific tools
+- Audit logging of tool usage
+- Input validation and sanitization
+- Output filtering and transformation
 
 Planned hooks (not yet implemented):
 
-- `tool_pre_invoke` / `tool_post_invoke` - Tool execution guardrails
 - `resource_pre_fetch` / `resource_post_fetch` - Resource content filtering
 - `server_pre_register` / `server_post_register` - Server validation
 - `auth_pre_check` / `auth_post_check` - Custom authentication
@@ -179,12 +194,16 @@ Planned hooks (not yet implemented):
 ```python
 from mcpgateway.plugins.framework.base import Plugin
 from mcpgateway.plugins.framework.models import PluginConfig
-from mcpgateway.plugins.framework.types import (
+from mcpgateway.plugins.framework.plugin_types import (
     PluginContext,
     PromptPrehookPayload,
     PromptPrehookResult,
     PromptPosthookPayload,
-    PromptPosthookResult
+    PromptPosthookResult,
+    ToolPreInvokePayload,
+    ToolPreInvokeResult,
+    ToolPostInvokePayload,
+    ToolPostInvokeResult
 )
 
 class MyPlugin(Plugin):
@@ -251,6 +270,62 @@ class MyPlugin(Plugin):
             modified_payload=payload
         )
 
+    async def tool_pre_invoke(
+        self,
+        payload: ToolPreInvokePayload,
+        context: PluginContext
+    ) -> ToolPreInvokeResult:
+        """Process tool before invocation."""
+
+        # Access tool name and arguments
+        tool_name = payload.name
+        args = payload.args
+
+        # Example: Block dangerous operations
+        if tool_name == "file_delete" and "system" in str(args):
+            return ToolPreInvokeResult(
+                continue_processing=False,
+                violation=PluginViolation(
+                    plugin_name=self.name,
+                    description="Dangerous operation blocked",
+                    violation_code="DANGEROUS_OP",
+                    details={"tool": tool_name}
+                )
+            )
+
+        # Example: Modify arguments
+        if "sanitize_me" in args:
+            args["sanitize_me"] = self.sanitize_input(args["sanitize_me"])
+            return ToolPreInvokeResult(
+                modified_payload=ToolPreInvokePayload(tool_name, args)
+            )
+
+        return ToolPreInvokeResult()
+
+    async def tool_post_invoke(
+        self,
+        payload: ToolPostInvokePayload,
+        context: PluginContext
+    ) -> ToolPostInvokeResult:
+        """Process tool after invocation."""
+
+        # Access tool result
+        tool_name = payload.name
+        result = payload.result
+
+        # Example: Filter sensitive data from results
+        if isinstance(result, dict) and "sensitive_data" in result:
+            result["sensitive_data"] = "[REDACTED]"
+            return ToolPostInvokeResult(
+                modified_payload=ToolPostInvokePayload(tool_name, result)
+            )
+
+        # Example: Add audit metadata
+        context.metadata["tool_executed"] = tool_name
+        context.metadata["execution_time"] = time.time()
+
+        return ToolPostInvokeResult()
+
     async def shutdown(self):
         """Cleanup when plugin shuts down."""
         # Close connections, save state, etc.

mcpgateway/db.py

@@ -1117,9 +1117,6 @@ class Gateway(Base):
     # Header passthrough configuration
     passthrough_headers: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)  # Store list of strings as JSON array
 
-    # Header passthrough configuration
-    passthrough_headers: Mapped[Optional[List[str]]] = mapped_column(JSON, nullable=True)  # Store list of strings as JSON array
-
     # Relationship with local tools this gateway provides
     tools: Mapped[List["Tool"]] = relationship(back_populates="gateway", foreign_keys="Tool.gateway_id", cascade="all, delete-orphan")
 

mcpgateway/plugins/framework/base.py

@@ -27,6 +27,10 @@
     PromptPosthookResult,
     PromptPrehookPayload,
     PromptPrehookResult,
+    ToolPostInvokePayload,
+    ToolPostInvokeResult,
+    ToolPreInvokePayload,
+    ToolPreInvokeResult,
 )
 
 
@@ -166,6 +170,46 @@ async def prompt_post_fetch(self, payload: PromptPosthookPayload, context: Plugi
                                    """
         )
 
+    async def tool_pre_invoke(self, payload: ToolPreInvokePayload, context: PluginContext) -> ToolPreInvokeResult:
+        """Plugin hook run before a tool is invoked.
+
+        Args:
+            payload: The tool payload to be analyzed.
+            context: Contextual information about the hook call.
+
+        Returns:
+            ToolPreInvokeResult with processing status and modified payload.
+
+        Examples:
+            >>> from mcpgateway.plugins.framework.plugin_types import ToolPreInvokePayload, PluginContext, GlobalContext
+            >>> payload = ToolPreInvokePayload("calculator", {"operation": "add", "a": 5, "b": 3})
+            >>> context = PluginContext(GlobalContext(request_id="123"))
+            >>> # In async context:
+            >>> # result = await plugin.tool_pre_invoke(payload, context)
+        """
+        # Default pass-through implementation
+        return ToolPreInvokeResult(continue_processing=True, modified_payload=payload)
+
+    async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
+        """Plugin hook run after a tool is invoked.
+
+        Args:
+            payload: The tool result payload to be analyzed.
+            context: Contextual information about the hook call.
+
+        Returns:
+            ToolPostInvokeResult with processing status and modified result.
+
+        Examples:
+            >>> from mcpgateway.plugins.framework.plugin_types import ToolPostInvokePayload, PluginContext, GlobalContext
+            >>> payload = ToolPostInvokePayload("calculator", {"result": 8, "status": "success"})
+            >>> context = PluginContext(GlobalContext(request_id="123"))
+            >>> # In async context:
+            >>> # result = await plugin.tool_post_invoke(payload, context)
+        """
+        # Default pass-through implementation
+        return ToolPostInvokeResult(continue_processing=True, modified_payload=payload)
+
     async def shutdown(self) -> None:
         """Plugin cleanup code."""
 

mcpgateway/plugins/framework/manager.py

@@ -46,9 +46,13 @@
     PromptPosthookResult,
     PromptPrehookPayload,
     PromptPrehookResult,
+    ToolPostInvokePayload,
+    ToolPostInvokeResult,
+    ToolPreInvokePayload,
+    ToolPreInvokeResult,
 )
 from mcpgateway.plugins.framework.registry import PluginInstanceRegistry
-from mcpgateway.plugins.framework.utils import post_prompt_matches, pre_prompt_matches
+from mcpgateway.plugins.framework.utils import post_prompt_matches, post_tool_matches, pre_prompt_matches, pre_tool_matches
 
 logger = logging.getLogger(__name__)
 
@@ -303,6 +307,54 @@ async def post_prompt_fetch(plugin: PluginRef, payload: PromptPosthookPayload, c
     return await plugin.plugin.prompt_post_fetch(payload, context)
 
 
+async def pre_tool_invoke(plugin: PluginRef, payload: ToolPreInvokePayload, context: PluginContext) -> ToolPreInvokeResult:
+    """Call plugin's tool pre-invoke hook.
+
+    Args:
+        plugin: The plugin to execute.
+        payload: The tool payload to be analyzed.
+        context: Contextual information about the hook call.
+
+    Returns:
+        The result of the plugin execution.
+
+    Examples:
+        >>> from mcpgateway.plugins.framework.base import Plugin, PluginRef
+        >>> from mcpgateway.plugins.framework.plugin_types import ToolPreInvokePayload, PluginContext, GlobalContext
+        >>> # Assuming you have a plugin instance:
+        >>> # plugin_ref = PluginRef(my_plugin)
+        >>> payload = ToolPreInvokePayload(name="calculator", args={"operation": "add", "a": 5, "b": 3})
+        >>> context = PluginContext(GlobalContext(request_id="123"))
+        >>> # In async context:
+        >>> # result = await pre_tool_invoke(plugin_ref, payload, context)
+    """
+    return await plugin.plugin.tool_pre_invoke(payload, context)
+
+
+async def post_tool_invoke(plugin: PluginRef, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
+    """Call plugin's tool post-invoke hook.
+
+    Args:
+        plugin: The plugin to execute.
+        payload: The tool result payload to be analyzed.
+        context: Contextual information about the hook call.
+
+    Returns:
+        The result of the plugin execution.
+
+    Examples:
+        >>> from mcpgateway.plugins.framework.base import Plugin, PluginRef
+        >>> from mcpgateway.plugins.framework.plugin_types import ToolPostInvokePayload, PluginContext, GlobalContext
+        >>> # Assuming you have a plugin instance:
+        >>> # plugin_ref = PluginRef(my_plugin)
+        >>> payload = ToolPostInvokePayload(name="calculator", result={"result": 8, "status": "success"})
+        >>> context = PluginContext(GlobalContext(request_id="123"))
+        >>> # In async context:
+        >>> # result = await post_tool_invoke(plugin_ref, payload, context)
+    """
+    return await plugin.plugin.tool_post_invoke(payload, context)
+
+
 class PluginManager:
     """Plugin manager for managing the plugin lifecycle.
 
@@ -343,6 +395,8 @@ class PluginManager:
     _config: Config | None = None
     _pre_prompt_executor: PluginExecutor[PromptPrehookPayload] = PluginExecutor[PromptPrehookPayload]()
     _post_prompt_executor: PluginExecutor[PromptPosthookPayload] = PluginExecutor[PromptPosthookPayload]()
+    _pre_tool_executor: PluginExecutor[ToolPreInvokePayload] = PluginExecutor[ToolPreInvokePayload]()
+    _post_tool_executor: PluginExecutor[ToolPostInvokePayload] = PluginExecutor[ToolPostInvokePayload]()
 
     # Context cleanup tracking
     _context_store: Dict[str, Tuple[PluginContextTable, float]] = {}
@@ -369,6 +423,8 @@ def __init__(self, config: str = "", timeout: int = DEFAULT_PLUGIN_TIMEOUT):
         # Update executor timeouts
         self._pre_prompt_executor.timeout = timeout
         self._post_prompt_executor.timeout = timeout
+        self._pre_tool_executor.timeout = timeout
+        self._post_tool_executor.timeout = timeout
 
         # Initialize context tracking if not already done
         if not hasattr(self, "_context_store"):
@@ -614,3 +670,112 @@ async def prompt_post_fetch(
             del self._context_store[global_context.request_id]
 
         return result
+
+    async def tool_pre_invoke(
+        self,
+        payload: ToolPreInvokePayload,
+        global_context: GlobalContext,
+        local_contexts: Optional[PluginContextTable] = None,
+    ) -> tuple[ToolPreInvokeResult, PluginContextTable | None]:
+        """Execute pre-invoke hooks before a tool is invoked.
+
+        Args:
+            payload: The tool payload containing name and arguments.
+            global_context: Shared context for all plugins with request metadata.
+            local_contexts: Optional existing contexts from previous executions.
+
+        Returns:
+            A tuple containing:
+            - ToolPreInvokeResult with processing status and modified payload
+            - PluginContextTable with updated contexts for post-invoke hook
+
+        Raises:
+            PayloadSizeError: If payload exceeds size limits.
+
+        Examples:
+            >>> manager = PluginManager("plugins/config.yaml")
+            >>> # In async context:
+            >>> # await manager.initialize()
+            >>>
+            >>> from mcpgateway.plugins.framework.plugin_types import ToolPreInvokePayload, GlobalContext
+            >>> payload = ToolPreInvokePayload(
+            ...     name="calculator",
+            ...     args={"operation": "add", "a": 5, "b": 3}
+            ... )
+            >>> context = GlobalContext(
+            ...     request_id="req-123",
+            ...     user="[email protected]"
+            ... )
+            >>>
+            >>> # In async context:
+            >>> # result, contexts = await manager.tool_pre_invoke(payload, context)
+            >>> # if result.continue_processing:
+            >>> #     # Proceed with tool invocation
+            >>> #     modified_payload = result.modified_payload or payload
+        """
+        # Cleanup old contexts periodically
+        await self._cleanup_old_contexts()
+
+        # Get plugins configured for this hook
+        plugins = self._registry.get_plugins_for_hook(HookType.TOOL_PRE_INVOKE)
+
+        # Execute plugins
+        result = await self._pre_tool_executor.execute(plugins, payload, global_context, pre_tool_invoke, pre_tool_matches, local_contexts)
+
+        # Store contexts for potential reuse
+        if result[1]:
+            self._context_store[global_context.request_id] = (result[1], time.time())
+
+        return result
+
+    async def tool_post_invoke(
+        self, payload: ToolPostInvokePayload, global_context: GlobalContext, local_contexts: Optional[PluginContextTable] = None
+    ) -> tuple[ToolPostInvokeResult, PluginContextTable | None]:
+        """Execute post-invoke hooks after a tool is invoked.
+
+        Args:
+            payload: The tool result payload containing invocation results.
+            global_context: Shared context for all plugins with request metadata.
+            local_contexts: Optional contexts from pre-invoke hook execution.
+
+        Returns:
+            A tuple containing:
+            - ToolPostInvokeResult with processing status and modified result
+            - PluginContextTable with final contexts
+
+        Raises:
+            PayloadSizeError: If payload exceeds size limits.
+
+        Examples:
+            >>> # Continuing from tool_pre_invoke example
+            >>> from mcpgateway.plugins.framework.plugin_types import ToolPostInvokePayload, GlobalContext
+            >>>
+            >>> post_payload = ToolPostInvokePayload(
+            ...     name="calculator",
+            ...     result={"result": 8, "status": "success"}
+            ... )
+            >>>
+            >>> manager = PluginManager("plugins/config.yaml")
+            >>> context = GlobalContext(request_id="req-123")
+            >>>
+            >>> # In async context:
+            >>> # result, _ = await manager.tool_post_invoke(
+            >>> #     post_payload,
+            >>> #     context,
+            >>> #     contexts  # From pre_invoke
+            >>> # )
+            >>> # if result.modified_payload:
+            >>> #     # Use modified result
+            >>> #     final_result = result.modified_payload.result
+        """
+        # Get plugins configured for this hook
+        plugins = self._registry.get_plugins_for_hook(HookType.TOOL_POST_INVOKE)
+
+        # Execute plugins
+        result = await self._post_tool_executor.execute(plugins, payload, global_context, post_tool_invoke, post_tool_matches, local_contexts)
+
+        # Clean up stored context after post-invoke
+        if global_context.request_id in self._context_store:
+            del self._context_store[global_context.request_id]
+
+        return result