Add support for custom messages

Author: cmaglie

client.go

@@ -41,14 +41,18 @@ type ServerMessagesHandler interface {
 
 // Client is an LSP Client
 type Client struct {
-	conn         *jsonrpc.Connection
-	handler      ServerMessagesHandler
-	errorHandler func(e error)
+	conn               *jsonrpc.Connection
+	handler            ServerMessagesHandler
+	customNotification map[string]CustomNotification
+	customRequest      map[string]CustomRequest
+	errorHandler       func(e error)
 }
 
 func NewClient(in io.Reader, out io.Writer, handler ServerMessagesHandler) *Client {
 	client := &Client{
-		errorHandler: func(e error) {},
+		errorHandler:       func(e error) {},
+		customNotification: map[string]CustomNotification{},
+		customRequest:      map[string]CustomRequest{},
 	}
 	client.handler = handler
 	client.conn = jsonrpc.NewConnection(
@@ -67,6 +71,14 @@ func (client *Client) SetErrorHandler(handler func(e error)) {
 	client.errorHandler = handler
 }
 
+func (client *Client) RegisterCustomNotification(method string, callback CustomNotification) {
+	client.customNotification[method] = callback
+}
+
+func (client *Client) RegisterCustomRequest(method string, callback CustomRequest) {
+	client.customRequest[method] = callback
+}
+
 func (client *Client) Run() {
 	client.conn.Run()
 }
@@ -114,7 +126,11 @@ func (client *Client) notificationDispatcher(logger jsonrpc.FunctionLogger, meth
 		}
 		client.handler.TextDocumentPublishDiagnostics(logger, &param)
 	default:
-		panic("unimplemented message")
+		if handler, ok := client.customNotification[method]; ok {
+			handler(logger, req)
+		} else {
+			panic("unimplemented notification: " + method)
+		}
 	}
 }
 
@@ -177,7 +193,11 @@ func (client *Client) requestDispatcher(ctx context.Context, logger jsonrpc.Func
 	case "workspace/codeLens/refresh":
 		resp(nil, client.handler.WorkspaceCodeLensRefresh(ctx, logger))
 	default:
-		panic("unimplemented message")
+		if handler, ok := client.customRequest[method]; ok {
+			resp(handler(ctx, logger, req))
+		} else {
+			panic("unimplemented request: " + method)
+		}
 	}
 }
 

server.go

@@ -86,14 +86,24 @@ type ClientMessagesHandler interface {
 
 // Server is an LSP Server
 type Server struct {
-	conn         *jsonrpc.Connection
-	handler      ClientMessagesHandler
-	errorHandler func(e error)
+	conn               *jsonrpc.Connection
+	handler            ClientMessagesHandler
+	customNotification map[string]CustomNotification
+	customRequest      map[string]CustomRequest
+	errorHandler       func(e error)
 }
 
+// CustomNotification is a function type for incoming custom notifications callbacks
+type CustomNotification func(logger jsonrpc.FunctionLogger, req json.RawMessage)
+
+// CustomNotification is a function type for incoming custom requests callbacks
+type CustomRequest func(ctx context.Context, logger jsonrpc.FunctionLogger, req json.RawMessage) (res interface{}, err *jsonrpc.ResponseError)
+
 func NewServer(in io.Reader, out io.Writer, handler ClientMessagesHandler) *Server {
 	serv := &Server{
-		errorHandler: func(e error) {},
+		errorHandler:       func(e error) {},
+		customNotification: map[string]CustomNotification{},
+		customRequest:      map[string]CustomRequest{},
 	}
 	serv.handler = handler
 	serv.conn = jsonrpc.NewConnection(
@@ -112,6 +122,14 @@ func (serv *Server) SetErrorHandler(handler func(e error)) {
 	serv.errorHandler = handler
 }
 
+func (serv *Server) RegisterCustomNotification(method string, callback CustomNotification) {
+	serv.customNotification[method] = callback
+}
+
+func (serv *Server) RegisterCustomRequest(method string, callback CustomRequest) {
+	serv.customRequest[method] = callback
+}
+
 func (serv *Server) Run() {
 	serv.conn.Run()
 }
@@ -228,7 +246,11 @@ func (serv *Server) notificationDispatcher(logger jsonrpc.FunctionLogger, method
 		}
 		serv.handler.TextDocumentDidClose(logger, &param)
 	default:
-		panic("unimplemented message")
+		if handler, ok := serv.customNotification[method]; ok {
+			handler(logger, req)
+		} else {
+			panic("unimplemented notification: " + method)
+		}
 	}
 }
 
@@ -537,7 +559,11 @@ func (serv *Server) requestDispatcher(ctx context.Context, logger jsonrpc.Functi
 		}
 		resp(serv.handler.TextDocumentMoniker(ctx, logger, &param))
 	default:
-		panic("unimplemented message")
+		if handler, ok := serv.customRequest[method]; ok {
+			resp(handler(ctx, logger, req))
+		} else {
+			panic("unimplemented request: " + method)
+		}
 	}
 }