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

class AgentPreInvokePayload(BaseModel):
    agent_id: str                                    # Agent identifier (can be modified for routing)
    messages: List[Message]                          # Conversation messages (can be filtered/transformed)
    tools: Optional[List[str]] = None                # Available tools (can be restricted)
    headers: Optional[HttpHeaderPayload] = None      # HTTP headers
    model: Optional[str] = None                      # Model override
    system_prompt: Optional[str] = None              # System instructions override
    parameters: Optional[Dict[str, Any]] = None      # LLM parameters (temperature, max_tokens, etc.)

class AgentPostInvokePayload(BaseModel):
    agent_id: str                                    # Agent identifier
    messages: List[Message]                          # Response messages (can be filtered/transformed)
    tool_calls: Optional[List[Dict[str, Any]]] = None  # Tool invocations made by agent

from mcpgateway.plugins.framework import (
    Plugin,
    PluginContext,
    AgentPreInvokePayload,
    AgentPreInvokeResult,
    PluginViolation,
)

class AgentSafetyPlugin(Plugin):
    """Filter agent interactions for safety."""

    async def agent_pre_invoke(
        self,
        payload: AgentPreInvokePayload,
        context: PluginContext
    ) -> AgentPreInvokeResult:
        # Restrict dangerous tools
        if payload.tools:
            safe_tools = [t for t in payload.tools if t not in ["file_delete", "system_exec"]]
            if len(safe_tools) < len(payload.tools):
                payload.tools = safe_tools
                self.logger.info(f"Restricted tools for agent {payload.agent_id}")

        # Filter offensive content in messages
        for msg in payload.messages:
            if self._contains_offensive_content(msg.content):
                return AgentPreInvokeResult(
                    continue_processing=False,
                    violation=PluginViolation(
                        code="OFFENSIVE_CONTENT",
                        reason="Message contains offensive content",
                        description="Agent request blocked due to policy violation"
                    )
                )

        return AgentPreInvokeResult(
            modified_payload=payload,
            metadata={"safety_checked": True}
        )

from abc import ABC
from mcpgateway.plugins.framework import Plugin, PluginConfig

class MyPlugin(Plugin):
    """Your plugin must inherit from Plugin."""

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

from mcpgateway.plugins.framework import (
    Plugin,
    PluginContext,
    ToolPreInvokePayload,
    ToolPreInvokeResult,
)

class ContentFilterPlugin(Plugin):
    """Convention-based hook - method name matches hook type."""

    async def tool_pre_invoke(
        self,
        payload: ToolPreInvokePayload,
        context: PluginContext
    ) -> ToolPreInvokeResult:
        """This hook is automatically discovered by its name."""

        # Block dangerous operations
        if payload.name == "file_delete" and "system" in str(payload.args):
            from mcpgateway.plugins.framework import PluginViolation
            return ToolPreInvokeResult(
                continue_processing=False,
                violation=PluginViolation(
                    code="DANGEROUS_OP",
                    reason="Dangerous operation blocked",
                    description=f"Cannot delete system files"
                )
            )

        # Modify arguments
        modified_args = {**payload.args, "processed": True}
        modified_payload = ToolPreInvokePayload(
            name=payload.name,
            args=modified_args,
            headers=payload.headers
        )

        return ToolPreInvokeResult(
            modified_payload=modified_payload,
            metadata={"processed_by": self.name}
        )

from mcpgateway.plugins.framework import Plugin, PluginContext
from mcpgateway.plugins.framework.decorator import hook
from mcpgateway.plugins.framework import (
    ToolHookType,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
)

class AuditPlugin(Plugin):
    """Decorator-based hook with descriptive method name."""

    @hook(ToolHookType.TOOL_POST_INVOKE)
    async def audit_tool_execution(
        self,
        payload: ToolPostInvokePayload,
        context: PluginContext
    ) -> ToolPostInvokeResult:
        """Method name doesn't match hook type, but @hook decorator registers it."""

        # Log tool execution
        self.logger.info(f"Tool executed: {payload.name}")

        # Filter sensitive data from results
        if isinstance(payload.result, dict) and "password" in payload.result:
            filtered_result = {**payload.result, "password": "[REDACTED]"}
            modified_payload = ToolPostInvokePayload(
                name=payload.name,
                result=filtered_result
            )
            return ToolPostInvokeResult(modified_payload=modified_payload)

        return ToolPostInvokeResult(continue_processing=True)

from mcpgateway.plugins.framework import (
    Plugin,
    PluginContext,
    PluginPayload,
    PluginResult
)
from mcpgateway.plugins.framework.decorator import hook

# Define custom payload type
class EmailPayload(PluginPayload):
    recipient: str
    subject: str
    body: str

# Define custom result type
class EmailResult(PluginResult[EmailPayload]):
    pass

class EmailPlugin(Plugin):
    """Custom hook with new hook type."""

    @hook("email_pre_send", EmailPayload, EmailResult)
    async def validate_email(
        self,
        payload: EmailPayload,
        context: PluginContext
    ) -> EmailResult:
        """Completely new hook type: 'email_pre_send'"""

        # Validate email address
        if "@" not in payload.recipient:
            modified_payload = EmailPayload(
                recipient=f"{payload.recipient}@example.com",
                subject=payload.subject,
                body=payload.body
            )
            return EmailResult(
                modified_payload=modified_payload,
                metadata={"fixed_email": True}
            )

        return EmailResult(continue_processing=True)

async def hook_name(
    self,
    payload: PayloadType,           # Specific to the hook (e.g., ToolPreInvokePayload)
    context: PluginContext          # Always PluginContext
) -> PluginResult[PayloadType]:     # PluginResult parameterized by payload type
    """Hook implementation."""
    pass

# These are type aliases defined in the framework
ToolPreInvokeResult = PluginResult[ToolPreInvokePayload]
ToolPostInvokeResult = PluginResult[ToolPostInvokePayload]
PromptPrehookResult = PluginResult[PromptPrehookPayload]
HttpAuthResolveUserResult = PluginResult[dict]  # Special case for user dict
# ... and so on for each hook type

# All of these are valid ways to construct results:
return ToolPreInvokeResult(continue_processing=True)
return ToolPreInvokeResult(modified_payload=new_payload)
return ToolPreInvokeResult(
    modified_payload=new_payload,
    metadata={"processed": True}
)

class MyPlugin(Plugin):
    async def initialize(self):
        """Called when plugin is loaded."""
        # Set up resources, connections, etc.
        self._session = aiohttp.ClientSession()

    async def shutdown(self):
        """Called when plugin manager shuts down."""
        # Cleanup resources
        if hasattr(self, '_session'):
            await self._session.close()

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