IBM Context Forge MCP plugin for dumping tool invocation context information.

plugin.py

"""
This IBM mcp-context-forge plugin logs all tool invocations (pre and post) with full context information.
"""

# Standard
import json
from typing import Any

# First-Party
from mcpgateway.plugins.framework import (
    Plugin,
    PluginConfig,
    PluginContext,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
    ToolPreInvokePayload,
    ToolPreInvokeResult,
)
from mcpgateway.plugins.framework.constants import TOOL_METADATA
from mcpgateway.services.logging_service import LoggingService

# Initialize logging service
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)


class ToolCallLoggerPlugin(Plugin):
    """Plugin that logs all tool invocations with full context information."""

    def __init__(self, config: PluginConfig):
        """Initialize the tool call logger plugin.

        Args:
            config: Plugin configuration.
        """
        super().__init__(config)

    def _log_tool_call(self, hook_type: str, payload: Any, context: PluginContext) -> None:
        """Log tool call information in a structured format.

        Args:
            hook_type: Type of hook (tool_pre_invoke or tool_post_invoke).
            payload: The tool payload (ToolPreInvokePayload or ToolPostInvokePayload).
            context: Plugin context with global context and metadata.
        """
        # Extract tool metadata if available
        tool_metadata = None
        if context.global_context.metadata and TOOL_METADATA in context.global_context.metadata:
            tool_metadata = context.global_context.metadata[TOOL_METADATA]

        # Build structured log entry
        log_entry = {
            "hook_type": hook_type,
            "tool_name": payload.name,
            "request_id": context.global_context.request_id,
            "tenant_id": context.global_context.tenant_id,
            "server_id": context.global_context.server_id,
            "user": context.global_context.user,
        }

        # Add tool-specific information
        if hook_type == "tool_pre_invoke":
            log_entry["args"] = payload.args if payload.args else {}
            if payload.headers:
                log_entry["headers"] = payload.headers.model_dump()
            else:
                log_entry["headers"] = {}
        elif hook_type == "tool_post_invoke":
            # Truncate large results to avoid log bloat
            result_str = str(payload.result)
            if len(result_str) > 1000:
                log_entry["result"] = result_str[:1000] + "... (truncated)"
                log_entry["result_length"] = len(result_str)
            else:
                log_entry["result"] = payload.result

        # Add tool metadata if available
        if tool_metadata:
            log_entry["tool_metadata"] = {
                "original_name": getattr(tool_metadata, "original_name", None),
                "url": str(getattr(tool_metadata, "url", None)) if hasattr(tool_metadata, "url") else None,
                "integration_type": getattr(tool_metadata, "integration_type", None),
            }

        # Log as JSON for easy parsing
        logger.info("Tool call logged: %s", json.dumps(log_entry, default=str))

    async def tool_pre_invoke(self, payload: ToolPreInvokePayload, context: PluginContext) -> ToolPreInvokeResult:
        """Plugin hook run before a tool is invoked.

        Logs tool name, arguments, headers, and context information.

        Args:
            payload: The tool payload to be analyzed.
            context: Contextual information about the hook call.

        Returns:
            The result allowing the tool to proceed.
        """
        self._log_tool_call("tool_pre_invoke", payload, context)
        return ToolPreInvokeResult(continue_processing=True)

    async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
        """Plugin hook run after a tool is invoked.

        Logs tool name, result, and context information.

        Args:
            payload: The tool result payload to be analyzed.
            context: Contextual information about the hook call.

        Returns:
            The result allowing the tool result to proceed.
        """
        self._log_tool_call("tool_post_invoke", payload, context)
        return ToolPostInvokeResult(continue_processing=True)