Plugin Framework¶

class Plugin:
    async def prompt_pre_fetch(self, payload, context) -> PluginResult
    async def tool_pre_invoke(self, payload, context) -> PluginResult
    # ... unified interface for all hook points

from mcpgateway.plugins.framework import (
    Plugin,
    PluginConfig,
    PluginContext,
    PromptPrehookPayload,
    PromptPrehookResult,
    PromptPosthookPayload,
    PromptPosthookResult,
    ToolPreInvokePayload,
    ToolPreInvokeResult,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
    ResourcePreFetchPayload,
    ResourcePreFetchResult,
    ResourcePostFetchPayload,
    ResourcePostFetchResult
)

class MyPlugin(Plugin):
    """Example plugin implementation."""

    def __init__(self, config: PluginConfig):
        super().__init__(config)
        # Initialize plugin-specific configuration
        self.my_setting = config.config.get("my_setting", "default")

    async def prompt_pre_fetch(
        self,
        payload: PromptPrehookPayload,
        context: PluginContext
    ) -> PromptPrehookResult:
        """Process prompt before retrieval."""

        # Access prompt name and arguments
        prompt_name = payload.name
        args = payload.args

        # Example: Block requests with forbidden words
        if "forbidden" in str(args.values()).lower():
            return PromptPrehookResult(
                continue_processing=False,
                violation=PluginViolation(
                    plugin_name=self.name,
                    description="Forbidden content detected",
                    violation_code="FORBIDDEN_CONTENT",
                    details={"found_in": "arguments"}
                )
            )

        # Example: Modify arguments
        if "transform_me" in args:
            args["transform_me"] = args["transform_me"].upper()
            return PromptPrehookResult(
                modified_payload=PromptPrehookPayload(prompt_name, args)
            )

        # Allow request to continue unmodified
        return PromptPrehookResult()

    async def prompt_post_fetch(
        self,
        payload: PromptPosthookPayload,
        context: PluginContext
    ) -> PromptPosthookResult:
        """Process prompt after rendering."""

        # Access rendered prompt
        prompt_result = payload.result

        # Example: Add metadata to context
        context.metadata["processed_by"] = self.name

        # Example: Modify response
        for message in prompt_result.messages:
            message.content.text = message.content.text.replace(
                "old_text", "new_text"
            )

        return PromptPosthookResult(
            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 resource_pre_fetch(
        self,
        payload: ResourcePreFetchPayload,
        context: PluginContext
    ) -> ResourcePreFetchResult:
        """Process resource before fetching."""

        # Access resource URI and metadata
        uri = payload.uri
        metadata = payload.metadata

        # Example: Block certain protocols
        from urllib.parse import urlparse
        parsed = urlparse(uri)
        if parsed.scheme not in ["http", "https", "file"]:
            return ResourcePreFetchResult(
                continue_processing=False,
                violation=PluginViolation(
                    plugin_name=self.name,
                    description=f"Protocol {parsed.scheme} not allowed",
                    violation_code="PROTOCOL_BLOCKED",
                    details={"uri": uri, "protocol": parsed.scheme}
                )
            )

        # Example: Add metadata
        metadata["validated_by"] = self.name
        return ResourcePreFetchResult(
            modified_payload=ResourcePreFetchPayload(uri, metadata)
        )

    async def resource_post_fetch(
        self,
        payload: ResourcePostFetchPayload,
        context: PluginContext
    ) -> ResourcePostFetchResult:
        """Process resource after fetching."""

        # Access resource content
        uri = payload.uri
        content = payload.content

        # Example: Redact sensitive patterns from text content
        if hasattr(content, 'text') and content.text:
            # Redact email addresses
            import re
            content.text = re.sub(
                r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
                '[EMAIL_REDACTED]',
                content.text
            )

        return ResourcePostFetchResult(
            modified_payload=ResourcePostFetchPayload(uri, content)
        )

    async def shutdown(self):
        """Cleanup when plugin shuts down."""
        # Close connections, save state, etc.
        pass

async def prompt_pre_fetch(self, payload, context):
    # Store state for later use
    context.set_state("request_time", time.time())
    context.set_state("original_args", payload.args.copy())

    return PromptPrehookResult()

async def prompt_post_fetch(self, payload, context):
    # Retrieve state from pre-hook
    elapsed = time.time() - context.get_state("request_time", 0)
    original = context.get_state("original_args", {})

    # Add timing metadata
    context.metadata["processing_time_ms"] = elapsed * 1000

    return PromptPosthookResult()

class LLMGuardPlugin(Plugin):
    """Example external service integration."""

    def __init__(self, config: PluginConfig):
        super().__init__(config)
        self.service_url = config.config.get("service_url")
        self.api_key = config.config.get("api_key")
        self.timeout = config.config.get("timeout", 30)

    async def prompt_pre_fetch(self, payload, context):
        # Call external service
        async with httpx.AsyncClient() as client:
            try:
                response = await client.post(
                    f"{self.service_url}/analyze",
                    json={
                        "text": str(payload.args),
                        "policy": "strict"
                    },
                    headers={
                        "Authorization": f"Bearer {self.api_key}"
                    },
                    timeout=self.timeout
                )

                result = response.json()

                if result.get("blocked", False):
                    return PromptPrehookResult(
                        continue_processing=False,
                        violation=PluginViolation(
                            plugin_name=self.name,
                            description=result.get("reason", "Content blocked"),
                            violation_code="LLMGUARD_BLOCKED",
                            details=result
                        )
                    )

            except Exception as e:
                # Handle errors based on plugin settings
                if self.config.mode == PluginMode.ENFORCE:
                    return PromptPrehookResult(
                        continue_processing=False,
                        violation=PluginViolation(
                            plugin_name=self.name,
                            description=f"Service error: {str(e)}",
                            violation_code="SERVICE_ERROR",
                            details={"error": str(e)}
                        )
                    )

        return PromptPrehookResult()

# plugins/my_plugin/plugin.py
from mcpgateway.plugins.framework import Plugin

class MyPlugin(Plugin):
    # Implementation here
    pass

# tests/test_my_plugin.py
import pytest
from plugins.my_plugin.plugin import MyPlugin
from mcpgateway.plugins.framework import PluginConfig

@pytest.mark.asyncio
async def test_my_plugin():
    config = PluginConfig(
        name="test",
        kind="plugins.my_plugin.plugin.MyPlugin",
        hooks=["prompt_pre_fetch"],
        config={"setting_one": "test_value"}
    )

    plugin = MyPlugin(config)

    # Test your plugin logic
    result = await plugin.prompt_pre_fetch(payload, context)
    assert result.continue_processing

async def prompt_pre_fetch(self, payload, context):
    try:
        # Plugin logic
        pass
    except Exception as e:
        logger.error(f"Plugin {self.name} error: {e}")

        # In permissive mode, log and continue
        if self.mode == PluginMode.PERMISSIVE:
            return PromptPrehookResult()

        # In enforce mode, block the request
        return PromptPrehookResult(
            continue_processing=False,
            violation=PluginViolation(
                plugin_name=self.name,
                description="Plugin error occurred",
                violation_code="PLUGIN_ERROR",
                details={"error": str(e)}
            )
        )

class CachedPlugin(Plugin):
    def __init__(self, config):
        super().__init__(config)
        self._cache = {}
        self._cache_ttl = config.config.get("cache_ttl", 300)

    async def expensive_operation(self, key):
        # Check cache first
        if key in self._cache:
            cached_value, timestamp = self._cache[key]
            if time.time() - timestamp < self._cache_ttl:
                return cached_value

        # Perform expensive operation
        result = await self._do_expensive_work(key)

        # Cache result
        self._cache[key] = (result, time.time())
        return result

logger.debug(f"Plugin {self.name} processing prompt: {payload.name}")
logger.info(f"Plugin {self.name} blocked request: {violation_code}")
logger.warning(f"Plugin {self.name} timeout approaching")
logger.error(f"Plugin {self.name} failed: {error}")