feat(plugins): added prompt posthook functionality with executor, fixed some linting issues, updated example plugin with posthook.

Signed-off-by: Teryl Taylor <[email protected]>

Author: Teryl Taylor

mcpgateway/plugins/framework/loader/config.py

@@ -36,7 +36,7 @@ def load_config(config: str, use_jinja: bool = True) -> Config:
         with open(os.path.normpath(config), "r", encoding="utf-8") as file:
             template = file.read()
             if use_jinja:
-                jinja_env = jinja2.Environment(loader=jinja2.BaseLoader())
+                jinja_env = jinja2.Environment(loader=jinja2.BaseLoader(), autoescape=True)
                 rendered_template = jinja_env.from_string(template).render(env=os.environ)
             else:
                 rendered_template = template

mcpgateway/plugins/framework/manager.py

@@ -10,50 +10,164 @@
 
 # Standard
 import logging
-from typing import Optional
+from typing import Any, Callable, Coroutine, Generic, Optional, TypeVar
 
 # First-Party
+from mcpgateway.plugins.framework.base import PluginRef
 from mcpgateway.plugins.framework.loader.config import ConfigLoader
 from mcpgateway.plugins.framework.loader.plugin import PluginLoader
-from mcpgateway.plugins.framework.models import Config, HookType, PluginMode
+from mcpgateway.plugins.framework.models import Config, HookType, PluginCondition, PluginMode
 from mcpgateway.plugins.framework.registry import PluginInstanceRegistry
 from mcpgateway.plugins.framework.types import (
     GlobalContext,
     PluginContext,
     PluginContextTable,
+    PluginResult,
     PromptPosthookPayload,
     PromptPosthookResult,
     PromptPrehookPayload,
     PromptPrehookResult,
 )
-from mcpgateway.plugins.framework.utils import pre_prompt_matches
+from mcpgateway.plugins.framework.utils import post_prompt_matches, pre_prompt_matches
 
 logger = logging.getLogger(__name__)
 
+T = TypeVar('T')
+
+
+class PluginExecutor(Generic[T]):
+    """Executes a list of plugins."""
+    async def execute(
+        self,
+        plugins: list[PluginRef],
+        payload: T,
+        global_context: GlobalContext,
+        plugin_run: Callable[[PluginRef, T, PluginContext], Coroutine[Any, Any, PluginResult[T]]],
+        compare: Callable[[T, list[PluginCondition], GlobalContext], bool],
+        local_contexts: Optional[PluginContextTable] = None,
+    ) -> tuple[PluginResult[T] | None, PluginContextTable | None]:
+        """Execute a plugins hook run before a prompt is retrieved and rendered.
+
+        Args:
+            plugins: the list of plugins to execute.
+            payload: the payload to be analyzed.
+            global_context: contextual information for all plugins.
+            plugin_run: async function for executing plugin hook.
+            compare: function for comparing conditional information with context and payload 
+            local_contexts: context local to a single plugin.
+
+        Returns:
+            The result of the plugin's analysis, including whether the prompt can proceed.
+        """
+        if not plugins:
+            return (PluginResult[T](modified_payload=None), None)
+
+        res_local_contexts = {}
+        combined_metadata = {}
+        current_payload: T | None = None
+        for pluginref in plugins:
+            if not pluginref.conditions or not compare(payload, pluginref.conditions, global_context):
+                continue
+            local_context_key = global_context.request_id + pluginref.uuid
+            if local_contexts and local_context_key in local_contexts:
+                local_context = local_contexts[local_context_key]
+            else:
+                local_context = PluginContext(global_context)
+            res_local_contexts[local_context_key] = local_context
+            result = await plugin_run(pluginref, payload, local_context)
+
+            if result.metadata:
+                combined_metadata.update(result.metadata)
+
+            if result.modified_payload is not None:
+                current_payload = result.modified_payload
+
+            if not result.continue_processing:
+                # Check execution mode
+                if pluginref.plugin.mode == PluginMode.ENFORCE:
+                    return (PluginResult[T](continue_processing=False, modified_payload=current_payload, error=result.error, metadata=combined_metadata), None)
+                elif pluginref.plugin.mode == PluginMode.PERMISSIVE:
+                    logger.warning(f"Plugin {pluginref.plugin.name} would block (permissive mode): {result.error}")
+
+        return (PluginResult[T](continue_processing=True, modified_payload=current_payload, error=None, metadata=combined_metadata), res_local_contexts)
+
+
+async def pre_prompt_fetch(plugin: PluginRef, payload: PromptPrehookPayload, context: PluginContext) -> PromptPrehookResult:
+    """Call plugin's prompt pre-fetch hook.
+
+    Args:
+        plugin:  the plugin to execute.
+        payload: the prompt payload to be analyzed.
+        context: contextual information about the hook call. Including why it was called.
+
+    Returns:
+        The result of the plugin execution.
+    """
+    return await plugin.plugin.prompt_pre_fetch(payload, context)
+
+
+async def post_prompt_fetch(plugin: PluginRef, payload: PromptPosthookPayload, context: PluginContext) -> PromptPosthookResult:
+    """Call plugin's prompt post-fetch hook.
+
+    Args:
+        plugin:  the plugin to execute.
+        payload: the prompt payload to be analyzed.
+        context: contextual information about the hook call. Including why it was called.
+
+    Returns:
+        The result of the plugin execution.
+    """
+    return await plugin.plugin.prompt_post_fetch(payload, context)
+
 
 class PluginManager:
     """Plugin manager for managing the plugin lifecycle."""
 
-    def __init__(self, config: str):
+    __shared_state: dict[Any, Any] = {}
+    _loader: PluginLoader = PluginLoader()
+    _initialized: bool = False
+    _registry: PluginInstanceRegistry = PluginInstanceRegistry()
+    _config: Config | None = None
+    _pre_prompt_executor: PluginExecutor[PromptPrehookPayload] = PluginExecutor[PromptPrehookPayload]()
+    _post_prompt_executor: PluginExecutor[PromptPosthookPayload] = PluginExecutor[PromptPosthookPayload]()
+
+    def __init__(self, config: str = ""):
         """Initialize plugin manager.
 
         Args:
             config: plugin configuration path.
         """
-        self._config: Config = ConfigLoader.load_config(config)
-        self._initialized: bool = False
-        self._loader: PluginLoader = PluginLoader()
-        self._registry: PluginInstanceRegistry = PluginInstanceRegistry()
+        self.__dict__ = self.__shared_state
+        if config:
+            self._config = ConfigLoader.load_config(config)
 
     @property
-    def config(self) -> Config:
+    def config(self) -> Config | None:
         """Plugin manager configuration.
 
         Returns:
             The plugin configuration.
         """
         return self._config
 
+    @property
+    def plugin_count(self) -> int:
+        """Number of plugins loaded.
+
+        Returns:
+            The number of plugins loaded.
+        """
+        return self._registry.plugin_count
+
+    @property
+    def initialized(self) -> bool:
+        """Plugin manager initialized.
+
+        Returns:
+            True if the plugin manager is initialized.
+        """
+        return self._initialized
+
     async def initialize(self) -> None:
         """Initialize the plugin manager.
 
@@ -62,8 +176,10 @@ async def initialize(self) -> None:
         """
         if self._initialized:
             return
+        
+        plugins = self._config.plugins if self._config else []
 
-        for plugin_config in self._config.plugins:
+        for plugin_config in plugins:
             if plugin_config.mode != PluginMode.DISABLED:
                 plugin = await self._loader.load_and_instantiate_plugin(plugin_config)
                 if plugin:
@@ -73,6 +189,16 @@ async def initialize(self) -> None:
         self._initialized = True
         logger.info(f"Plugin manager initialized with {len(self._registry.get_all_plugins())} plugins")
 
+    async def shutdown(self) -> None:
+        """Shutdown all plugins."""
+        for plugin_ref in self._registry.get_all_plugins():
+            try:
+                await plugin_ref.plugin.shutdown()
+            except Exception as e:
+                logger.error(f"Error shutting down plugin {plugin_ref.plugin.name}: {e}")
+
+        self._initialized = False
+
     async def prompt_pre_fetch(
         self,
         payload: PromptPrehookPayload,
@@ -90,38 +216,9 @@ async def prompt_pre_fetch(
             The result of the plugin's analysis, including whether the prompt can proceed.
         """
         plugins = self._registry.get_plugins_for_hook(HookType.PROMPT_PRE_FETCH)
+        return await self._pre_prompt_executor.execute(plugins, payload, global_context, pre_prompt_fetch, pre_prompt_matches, local_contexts)
 
-        if not plugins:
-            return (PromptPrehookResult(modified_payload=payload), None)
-
-        res_local_contexts = {}
-        combined_metadata = {}
-        current_payload: PromptPrehookPayload | None = None
-        for pluginref in plugins:
-            if not pluginref.conditions or not pre_prompt_matches(payload, pluginref.conditions, global_context):
-                continue
-            local_context_key = global_context.request_id + pluginref.uuid
-            if local_contexts and local_context_key in local_contexts:
-                local_context = local_contexts[local_context_key]
-            else:
-                local_context = PluginContext(global_context)
-            res_local_contexts[local_context_key] = local_context
-            result = await pluginref.plugin.prompt_pre_fetch(payload, local_context)
-
-            if result.metadata:
-                combined_metadata.update(result.metadata)
-
-            if result.modified_payload is not None:
-                current_payload = result.modified_payload
-
-            if not result.continue_processing:
-                # Check execution mode
-                if pluginref.plugin.mode == PluginMode.ENFORCE:
-                    return (PromptPrehookResult(continue_processing=False, modified_payload=current_payload, error=result.error, metadata=combined_metadata), None)
-                elif pluginref.plugin.mode == PluginMode.PERMISSIVE:
-                    logger.warning(f"Plugin {pluginref.plugin.name} would block (permissive mode): {result.error}")
 
-        return (PromptPrehookResult(continue_processing=True, modified_payload=current_payload, error=None, metadata=combined_metadata), res_local_contexts)
 
     async def prompt_post_fetch(
         self, payload: PromptPosthookPayload, global_context: GlobalContext, local_contexts: Optional[PluginContextTable] = None
@@ -136,4 +233,5 @@ async def prompt_post_fetch(
         Returns:
             The result of the plugin's analysis, including whether the prompt can proceed.
         """
-        return (None, None)
+        plugins = self._registry.get_plugins_for_hook(HookType.PROMPT_POST_FETCH)
+        return await self._post_prompt_executor.execute(plugins, payload, global_context, post_prompt_fetch, post_prompt_matches, local_contexts)

mcpgateway/plugins/framework/registry.py

@@ -105,3 +105,11 @@ def get_all_plugins(self) -> list[PluginRef]:
             A list of registered plugin instances.
         """
         return list(self._plugins.values())
+
+    def plugin_count(self) -> int:
+        """Return the number of plugins registered.
+
+        Returns:
+            The number of plugins registered.
+        """
+        return len(self._plugins)

mcpgateway/plugins/framework/types.py

@@ -30,10 +30,21 @@ def __init__(self, name: str, args: Optional[dict[str, str]]):
             args: The prompt arguments for rendering.
         """
         self.name = name
-        self.args = args
+        self.args = args or {}
 
 
-PromptPosthookPayload = PromptResult
+class PromptPosthookPayload:
+    """A prompt payload for a prompt posthook."""
+
+    def __init__(self, name: str, result: PromptResult):
+        """Initialize a prompt posthook payload.
+
+        Args:
+            name: The prompt name.
+            result: The prompt Prompt Result.
+        """
+        self.name = name
+        self.result = result
 
 
 class PluginResult(Generic[T]):
@@ -66,7 +77,7 @@ def __init__(
         request_id: str,
         user: Optional[str] = None,
         tenant_id: Optional[str] = None,
-        server_id: Optional[str] = None,
+        server_id: Optional[str] = None
     ) -> None:
         """Initialize a global context.
 
@@ -81,7 +92,6 @@ def __init__(
         self.tenant_id = tenant_id
         self.server_id = server_id
 
-
 class PluginContext(GlobalContext):
     """The plugin's context, which lasts a request lifecycle.
 

mcpgateway/plugins/framework/utils.py

@@ -16,7 +16,7 @@
 
 # First-Party
 from mcpgateway.plugins.framework.models import PluginCondition
-from mcpgateway.plugins.framework.types import GlobalContext, PromptPrehookPayload
+from mcpgateway.plugins.framework.types import GlobalContext, PromptPosthookPayload, PromptPrehookPayload
 
 
 @cache  # noqa
@@ -95,3 +95,27 @@ def pre_prompt_matches(payload: PromptPrehookPayload, conditions: list[PluginCon
         elif index < len(conditions) - 1:
             current_result = True
     return current_result
+
+def post_prompt_matches(payload: PromptPosthookPayload, conditions: list[PluginCondition], context: GlobalContext) -> bool:
+    """Check for a match on pre-prompt hooks.
+
+    Args:
+        payload: the prompt posthook payload.
+        conditions: the conditions on the plugin that are required for execution.
+        context: the global context.
+
+    Returns:
+        True if the plugin matches criteria.
+    """
+    current_result = True
+    for index, condition in enumerate(conditions):
+        if not matches(condition, context):
+            current_result = False
+
+        if condition.prompts and payload.name not in condition.prompts:
+            current_result = False
+        if current_result:
+            return True
+        elif index < len(conditions) - 1:
+            current_result = True
+    return current_result

plugins/config.yaml

@@ -20,6 +20,8 @@ plugins:
       words:
         - search: crap
           replace: crud
+        - search: crud
+          replace: yikes
 
 
 # Plugin directories to scan

plugins/regex/plugin-manifest.yaml

@@ -6,4 +6,4 @@ available_hooks:
   - "prompt_post_hook"
   - "tool_pre_hook"
   - "tool_post_hook"
-default_configs:
+default_configs: