Plugin Framework¶

class Plugin:
    async def prompt_pre_fetch(self, payload: PromptPrehookPayload,
                              context: PluginContext) -> PromptPrehookResult: ...
    async def prompt_post_fetch(self, payload: PromptPosthookPayload,
                               context: PluginContext) -> PromptPosthookResult: ...
    async def tool_pre_invoke(self, payload: ToolPreInvokePayload,
                             context: PluginContext) -> ToolPreInvokeResult: ...
    async def tool_post_invoke(self, payload: ToolPostInvokePayload,
                              context: PluginContext) -> ToolPostInvokeResult: ...
    async def resource_pre_fetch(self, payload: ResourcePreFetchPayload,
                                context: PluginContext) -> ResourcePreFetchResult: ...
    async def resource_post_fetch(self, payload: ResourcePostFetchPayload,
                                 context: PluginContext) -> ResourcePostFetchResult: ...
    # ... additional hook methods

from mcpgateway.plugins.framework import Plugin, PluginConfig, PluginContext, PromptPrehookPayload, PromptPrehookResult

class MyPlugin(Plugin):
    def __init__(self, config: PluginConfig):
        super().__init__(config)

    async def prompt_pre_fetch(self, payload: PromptPrehookPayload, context: PluginContext) -> PromptPrehookResult:
        # modify
        return PromptPrehookResult(modified_payload=payload)

        # or block
        # return PromptPrehookResult(
        #     continue_processing=False,
        #     violation=PluginViolation(
        #         reason=f"Blocked by {self.name}",
        #         description="...",
        #         code="...",
        #         details="..."
        #     )
        # )

class PromptPrehookPayload(BaseModel):
    name: str                                    # Prompt template name
    args: Optional[dict[str, str]] = Field(default_factory=dict)  # Template arguments

payload = PromptPrehookPayload(
    name="user_greeting",
    args={"user_name": "Alice", "time_of_day": "morning"}
)

class PromptPosthookPayload(BaseModel):
    name: str                                    # Prompt name
    result: PromptResult                         # Rendered prompt result

class ToolPreInvokePayload(BaseModel):
    name: str                                    # Tool name
    args: Optional[dict[str, Any]] = Field(default_factory=dict)  # Tool arguments
    headers: Optional[HttpHeaderPayload] = None  # HTTP pass-through headers

class ToolPostInvokePayload(BaseModel):
    name: str                                    # Tool name
    result: Any                                  # Tool execution result

class HttpHeaderPayload(RootModel[dict[str, str]]):
    # Provides dictionary-like access to HTTP headers
    # Supports: __iter__, __getitem__, __setitem__, __len__

headers = HttpHeaderPayload({"Authorization": "Bearer token", "Content-Type": "application/json"})
headers["X-Custom-Header"] = "custom_value"
auth_header = headers["Authorization"]

class ResourcePreFetchPayload(BaseModel):
    uri: str                                     # Resource URI
    metadata: Optional[dict[str, Any]] = Field(default_factory=dict)  # Request metadata

class ResourcePostFetchPayload(BaseModel):
    uri: str                                     # Resource URI
    content: Any                                 # Fetched resource content

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(
                    reason="Forbidden content",
                    description="Forbidden content detected",
                    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(
                    reason="Dangerous operation blocked",
                    description="Dangerous operation blocked",
                    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(
                reason="Protocol not allowed",
                description=f"Protocol {parsed.scheme} not allowed",
                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

class GlobalContext(BaseModel):
    """The global context, which shared across all plugins.

    Attributes:
            request_id (str): ID of the HTTP request.
            user (str): user ID associated with the request.
            tenant_id (str): tenant ID.
            server_id (str): server ID.
            metadata (Optional[dict[str,Any]]): a global shared metadata across plugins (Read-only from plugin's perspective.).
            state (Optional[dict[str,Any]]): a global shared state across plugins.
    """

    request_id: str
    user: Optional[str] = None
    tenant_id: Optional[str] = None
    server_id: Optional[str] = None
    state: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)


class PluginContext(BaseModel):
    """The plugin's context, which lasts a request lifecycle.

    Attributes:
       state:  the inmemory state of the request.
       global_context: the context that is shared across plugins.
       metadata: plugin meta data.
    """

    state: dict[str, Any] = Field(default_factory=dict)
    global_context: GlobalContext
    metadata: dict[str, Any] = Field(default_factory=dict)

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()

from mcpgateway.plugins.framework.constants import GATEWAY_METADATA, TOOL_METADATA

tool_meta = context.global_context.metadata[TOOL_METADATA]
assert tool_meta.original_name == "test_tool"
assert tool_meta.url.host == "example.com"
assert tool_meta.integration_type == "REST" or tool_meta.integration_type == "MCP"

gateway_meta = context.global_context.metadata[GATEWAY_METADATA]
assert gateway_meta.name == "test_gateway"
assert gateway_meta.transport == "sse"
assert gateway_meta.url.host == "example.com"

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(
                            reason="External service blocked",
                            description=result.get("reason", "Content blocked"),
                            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(
                            reason="Service error",
                            description=f"Service error: {str(e)}",
                            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

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}")