Class LoggingClientSession

optillm/plugins/mcp_plugin.py:232–259

A ClientSession that logs all communication

# Create a custom ClientSession that logs all communication
class LoggingClientSession(ClientSession):
    """A ClientSession that logs all communication"""
    
    async def send_request(self, *args, **kwargs):
        """Log and forward requests"""
        method = args[0]
        params = args[1] if len(args) > 1 else None
        log_mcp_message("REQUEST", method, params)
        
        try:
            result = await super().send_request(*args, **kwargs)
            log_mcp_message("RESPONSE", method, result=result)
            return result
        except Exception as e:
            log_mcp_message("ERROR", method, error=str(e))
            raise
    
    async def send_notification(self, *args, **kwargs):
        """Log and forward notifications"""
        method = args[0]
        params = args[1] if len(args) > 1 else None
        log_mcp_message("NOTIFICATION", method, params)
        
        try:
            await super().send_notification(*args, **kwargs)
        except Exception as e:
            log_mcp_message("ERROR", method, error=str(e))
            raise

class MCPServer:
    """Represents a connection to an MCP server"""