internal/jsonrpc2: refactor to enable a more advanced request

This separates hides the wire structures, and then exposes a new Request
type to allow for it to carry advanced features.
It also embeds the connection into the request and changes the signature of the
handler to no longer require a separate Conn argument.

Change-Id: I20b54f146285f7a9cb5f279c6ebdf0f286f4b829
Reviewed-on: https://go-review.googlesource.com/c/tools/+/183717
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/jsonrpc2/jsonrpc2.go b/internal/jsonrpc2/jsonrpc2.go
index db4a73b..a964113 100644
--- a/internal/jsonrpc2/jsonrpc2.go
+++ b/internal/jsonrpc2/jsonrpc2.go
@@ -33,14 +33,27 @@
 	stream             Stream
 	err                error
 	pendingMu          sync.Mutex // protects the pending map
-	pending            map[ID]chan *Response
+	pending            map[ID]chan *wireResponse
 	handlingMu         sync.Mutex // protects the handling map
 	handling           map[ID]handling
 }
 
+// Request is sent to a server to represent a Call or Notify operaton.
+type Request struct {
+	conn *Conn
+
+	// Method is a string containing the method name to invoke.
+	Method string
+	// Params is either a struct or an array with the parameters of the method.
+	Params *json.RawMessage
+	// The id of this request, used to tie the response back to the request.
+	// Will be either a string or a number. If not set, the request is a notify,
+	// and no response is possible.
+	ID *ID
+}
+
 type queueEntry struct {
 	ctx  context.Context
-	c    *Conn
 	r    *Request
 	size int64
 }
@@ -50,16 +63,14 @@
 // call Reply on the Conn with the supplied request.
 // Handlers are called synchronously, they should pass the work off to a go
 // routine if they are going to take a long time.
-type Handler func(context.Context, *Conn, *Request)
+type Handler func(context.Context, *Request)
 
 // Canceler is an option you can pass to NewConn which is invoked for
 // cancelled outgoing requests.
-// The request will have the ID filled in, which can be used to propagate the
-// cancel to the other process if needed.
 // It is okay to use the connection to send notifications, but the context will
 // be in the cancelled state, so you must do it with the background context
 // instead.
-type Canceler func(context.Context, *Conn, *Request)
+type Canceler func(context.Context, *Conn, ID)
 
 type rpcStats struct {
 	server   bool
@@ -133,17 +144,17 @@
 func NewConn(s Stream) *Conn {
 	conn := &Conn{
 		stream:   s,
-		pending:  make(map[ID]chan *Response),
+		pending:  make(map[ID]chan *wireResponse),
 		handling: make(map[ID]handling),
 	}
 	// the default handler reports a method error
-	conn.Handler = func(ctx context.Context, c *Conn, r *Request) {
+	conn.Handler = func(ctx context.Context, r *Request) {
 		if r.IsNotify() {
-			c.Reply(ctx, r, nil, NewErrorf(CodeMethodNotFound, "method %q not found", r.Method))
+			r.Reply(ctx, nil, NewErrorf(CodeMethodNotFound, "method %q not found", r.Method))
 		}
 	}
-	// the default canceller does nothing
-	conn.Canceler = func(context.Context, *Conn, *Request) {}
+	// the default canceler does nothing
+	conn.Canceler = func(context.Context, *Conn, ID) {}
 	// the default logger does nothing
 	conn.Logger = func(Direction, *ID, time.Duration, string, *json.RawMessage, *Error) {}
 	return conn
@@ -174,7 +185,7 @@
 	if err != nil {
 		return fmt.Errorf("marshalling notify parameters: %v", err)
 	}
-	request := &Request{
+	request := &wireRequest{
 		Method: method,
 		Params: jsonParams,
 	}
@@ -200,7 +211,7 @@
 	if err != nil {
 		return fmt.Errorf("marshalling call parameters: %v", err)
 	}
-	request := &Request{
+	request := &wireRequest{
 		ID:     &id,
 		Method: method,
 		Params: jsonParams,
@@ -212,7 +223,7 @@
 	}
 	// we have to add ourselves to the pending map before we send, otherwise we
 	// are racing the response
-	rchan := make(chan *Response)
+	rchan := make(chan *wireResponse)
 	c.pendingMu.Lock()
 	c.pending[id] = rchan
 	c.pendingMu.Unlock()
@@ -249,30 +260,38 @@
 		return nil
 	case <-ctx.Done():
 		// allow the handler to propagate the cancel
-		c.Canceler(ctx, c, request)
+		c.Canceler(ctx, c, id)
 		return ctx.Err()
 	}
 }
 
+// Conn returns the connection that created this request.
+func (r *Request) Conn() *Conn { return r.conn }
+
+// IsNotify returns true if this request is a notification.
+func (r *Request) IsNotify() bool {
+	return r.ID == nil
+}
+
 // Reply sends a reply to the given request.
 // It is an error to call this if request was not a call.
 // You must call this exactly once for any given request.
 // If err is set then result will be ignored.
-func (c *Conn) Reply(ctx context.Context, req *Request, result interface{}, err error) error {
-	ctx, st := trace.StartSpan(ctx, req.Method+":reply", trace.WithSpanKind(trace.SpanKindClient))
+func (r *Request) Reply(ctx context.Context, result interface{}, err error) error {
+	ctx, st := trace.StartSpan(ctx, r.Method+":reply", trace.WithSpanKind(trace.SpanKindClient))
 	defer st.End()
 
-	if req.IsNotify() {
+	if r.IsNotify() {
 		return fmt.Errorf("reply not invoked with a valid call")
 	}
-	c.handlingMu.Lock()
-	handling, found := c.handling[*req.ID]
+	r.conn.handlingMu.Lock()
+	handling, found := r.conn.handling[*r.ID]
 	if found {
-		delete(c.handling, *req.ID)
+		delete(r.conn.handling, *r.ID)
 	}
-	c.handlingMu.Unlock()
+	r.conn.handlingMu.Unlock()
 	if !found {
-		return fmt.Errorf("not a call in progress: %v", req.ID)
+		return fmt.Errorf("not a call in progress: %v", r.ID)
 	}
 
 	elapsed := time.Since(handling.start)
@@ -280,9 +299,9 @@
 	if err == nil {
 		raw, err = marshalToRaw(result)
 	}
-	response := &Response{
+	response := &wireResponse{
 		Result: raw,
-		ID:     req.ID,
+		ID:     r.ID,
 	}
 	if err != nil {
 		if callErr, ok := err.(*Error); ok {
@@ -295,8 +314,8 @@
 	if err != nil {
 		return err
 	}
-	c.Logger(Send, response.ID, elapsed, req.Method, response.Result, response.Error)
-	n, err := c.stream.Write(ctx, data)
+	r.conn.Logger(Send, response.ID, elapsed, r.Method, response.Result, response.Error)
+	n, err := r.conn.stream.Write(ctx, data)
 
 	v := ctx.Value(rpcStatsKey)
 	if v != nil {
@@ -332,7 +351,7 @@
 }
 
 func (c *Conn) deliver(ctx context.Context, q chan queueEntry, request *Request, size int64) bool {
-	e := queueEntry{ctx: ctx, c: c, r: request, size: size}
+	e := queueEntry{ctx: ctx, r: request, size: size}
 	if !c.RejectIfOverloaded {
 		q <- e
 		return true
@@ -361,7 +380,7 @@
 			}
 			ctx, rpcStats := start(ctx, true, e.r.Method, e.r.ID)
 			rpcStats.received += e.size
-			c.Handler(ctx, e.c, e.r)
+			c.Handler(ctx, e.r)
 			rpcStats.end(ctx, nil)
 		}
 	}()
@@ -385,6 +404,7 @@
 		case msg.Method != "":
 			// if method is set it must be a request
 			request := &Request{
+				conn:   c,
 				Method: msg.Method,
 				Params: msg.Params,
 				ID:     msg.ID,
@@ -407,7 +427,7 @@
 				c.Logger(Receive, request.ID, -1, request.Method, request.Params, nil)
 				if !c.deliver(reqCtx, q, request, n) {
 					// queue is full, reject the message by directly replying
-					c.Reply(ctx, request, nil, NewErrorf(CodeServerOverloaded, "no room in queue"))
+					request.Reply(ctx, nil, NewErrorf(CodeServerOverloaded, "no room in queue"))
 				}
 			}
 		case msg.ID != nil:
@@ -419,7 +439,7 @@
 			}
 			c.pendingMu.Unlock()
 			// and send the reply to the channel
-			response := &Response{
+			response := &wireResponse{
 				Result: msg.Result,
 				Error:  msg.Error,
 				ID:     msg.ID,
diff --git a/internal/jsonrpc2/jsonrpc2_test.go b/internal/jsonrpc2/jsonrpc2_test.go
index 5b4e0fb..c052f8c 100644
--- a/internal/jsonrpc2/jsonrpc2_test.go
+++ b/internal/jsonrpc2/jsonrpc2_test.go
@@ -122,36 +122,36 @@
 	return conn
 }
 
-func handle(ctx context.Context, c *jsonrpc2.Conn, r *jsonrpc2.Request) {
+func handle(ctx context.Context, r *jsonrpc2.Request) {
 	switch r.Method {
 	case "no_args":
 		if r.Params != nil {
-			c.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
+			r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
 			return
 		}
-		c.Reply(ctx, r, true, nil)
+		r.Reply(ctx, true, nil)
 	case "one_string":
 		var v string
 		if err := json.Unmarshal(*r.Params, &v); err != nil {
-			c.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
+			r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
 			return
 		}
-		c.Reply(ctx, r, "got:"+v, nil)
+		r.Reply(ctx, "got:"+v, nil)
 	case "one_number":
 		var v int
 		if err := json.Unmarshal(*r.Params, &v); err != nil {
-			c.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
+			r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
 			return
 		}
-		c.Reply(ctx, r, fmt.Sprintf("got:%d", v), nil)
+		r.Reply(ctx, fmt.Sprintf("got:%d", v), nil)
 	case "join":
 		var v []string
 		if err := json.Unmarshal(*r.Params, &v); err != nil {
-			c.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
+			r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err.Error()))
 			return
 		}
-		c.Reply(ctx, r, path.Join(v...), nil)
+		r.Reply(ctx, path.Join(v...), nil)
 	default:
-		c.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
+		r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
 	}
 }
diff --git a/internal/jsonrpc2/wire.go b/internal/jsonrpc2/wire.go
index bcf4d65..1e126a4 100644
--- a/internal/jsonrpc2/wire.go
+++ b/internal/jsonrpc2/wire.go
@@ -34,8 +34,8 @@
 	CodeServerOverloaded = -32000
 )
 
-// Request is sent to a server to represent a Call or Notify operaton.
-type Request struct {
+// wireRequest is sent to a server to represent a Call or Notify operaton.
+type wireRequest struct {
 	// VersionTag is always encoded as the string "2.0"
 	VersionTag VersionTag `json:"jsonrpc"`
 	// Method is a string containing the method name to invoke.
@@ -48,11 +48,11 @@
 	ID *ID `json:"id,omitempty"`
 }
 
-// Response is a reply to a Request.
+// wireResponse is a reply to a Request.
 // It will always have the ID field set to tie it back to a request, and will
 // have either the Result or Error fields set depending on whether it is a
 // success or failure response.
-type Response struct {
+type wireResponse struct {
 	// VersionTag is always encoded as the string "2.0"
 	VersionTag VersionTag `json:"jsonrpc"`
 	// Result is the response value, and is required on success.
@@ -87,11 +87,6 @@
 	Number int64
 }
 
-// IsNotify returns true if this request is a notification.
-func (r *Request) IsNotify() bool {
-	return r.ID == nil
-}
-
 func (err *Error) Error() string {
 	if err == nil {
 		return ""
diff --git a/internal/lsp/protocol/protocol.go b/internal/lsp/protocol/protocol.go
index c164e06..a956fe4 100644
--- a/internal/lsp/protocol/protocol.go
+++ b/internal/lsp/protocol/protocol.go
@@ -14,8 +14,8 @@
 const defaultMessageBufferSize = 20
 const defaultRejectIfOverloaded = false
 
-func canceller(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {
-	conn.Notify(context.Background(), "$/cancelRequest", &CancelParams{ID: *req.ID})
+func canceller(ctx context.Context, conn *jsonrpc2.Conn, id jsonrpc2.ID) {
+	conn.Notify(context.Background(), "$/cancelRequest", &CancelParams{ID: id})
 }
 
 func NewClient(stream jsonrpc2.Stream, client Client) (*jsonrpc2.Conn, Server, xlog.Logger) {
@@ -39,11 +39,11 @@
 	return conn, client, log
 }
 
-func sendParseError(ctx context.Context, log xlog.Logger, conn *jsonrpc2.Conn, req *jsonrpc2.Request, err error) {
+func sendParseError(ctx context.Context, log xlog.Logger, req *jsonrpc2.Request, err error) {
 	if _, ok := err.(*jsonrpc2.Error); !ok {
 		err = jsonrpc2.NewErrorf(jsonrpc2.CodeParseError, "%v", err)
 	}
-	if err := conn.Reply(ctx, req, nil, err); err != nil {
+	if err := req.Reply(ctx, nil, err); err != nil {
 		log.Errorf(ctx, "%v", err)
 	}
 }
diff --git a/internal/lsp/protocol/tsclient.go b/internal/lsp/protocol/tsclient.go
index f77d023..ccfa852 100644
--- a/internal/lsp/protocol/tsclient.go
+++ b/internal/lsp/protocol/tsclient.go
@@ -24,19 +24,19 @@
 }
 
 func clientHandler(log xlog.Logger, client Client) jsonrpc2.Handler {
-	return func(ctx context.Context, conn *jsonrpc2.Conn, r *jsonrpc2.Request) {
+	return func(ctx context.Context, r *jsonrpc2.Request) {
 		switch r.Method {
 		case "$/cancelRequest":
 			var params CancelParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
-			conn.Cancel(params.ID)
+			r.Conn().Cancel(params.ID)
 		case "window/showMessage": // notif
 			var params ShowMessageParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := client.ShowMessage(ctx, &params); err != nil {
@@ -45,7 +45,7 @@
 		case "window/logMessage": // notif
 			var params LogMessageParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := client.LogMessage(ctx, &params); err != nil {
@@ -54,7 +54,7 @@
 		case "telemetry/event": // notif
 			var params interface{}
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := client.Event(ctx, &params); err != nil {
@@ -63,7 +63,7 @@
 		case "textDocument/publishDiagnostics": // notif
 			var params PublishDiagnosticsParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := client.PublishDiagnostics(ctx, &params); err != nil {
@@ -71,67 +71,67 @@
 			}
 		case "workspace/workspaceFolders": // req
 			if r.Params != nil {
-				conn.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
+				r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
 				return
 			}
 			resp, err := client.WorkspaceFolders(ctx)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "workspace/configuration": // req
 			var params ConfigurationParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := client.Configuration(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "client/registerCapability": // req
 			var params RegistrationParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			err := client.RegisterCapability(ctx, &params)
-			if err := conn.Reply(ctx, r, nil, err); err != nil {
+			if err := r.Reply(ctx, nil, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "client/unregisterCapability": // req
 			var params UnregistrationParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			err := client.UnregisterCapability(ctx, &params)
-			if err := conn.Reply(ctx, r, nil, err); err != nil {
+			if err := r.Reply(ctx, nil, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "window/showMessageRequest": // req
 			var params ShowMessageRequestParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := client.ShowMessageRequest(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "workspace/applyEdit": // req
 			var params ApplyWorkspaceEditParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := client.ApplyEdit(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 
 		default:
 			if r.IsNotify() {
-				conn.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
+				r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
 			}
 		}
 	}
diff --git a/internal/lsp/protocol/tsserver.go b/internal/lsp/protocol/tsserver.go
index 607208b..8434861 100644
--- a/internal/lsp/protocol/tsserver.go
+++ b/internal/lsp/protocol/tsserver.go
@@ -55,19 +55,19 @@
 }
 
 func serverHandler(log xlog.Logger, server Server) jsonrpc2.Handler {
-	return func(ctx context.Context, conn *jsonrpc2.Conn, r *jsonrpc2.Request) {
+	return func(ctx context.Context, r *jsonrpc2.Request) {
 		switch r.Method {
 		case "$/cancelRequest":
 			var params CancelParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
-			conn.Cancel(params.ID)
+			r.Conn().Cancel(params.ID)
 		case "workspace/didChangeWorkspaceFolders": // notif
 			var params DidChangeWorkspaceFoldersParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidChangeWorkspaceFolders(ctx, &params); err != nil {
@@ -76,7 +76,7 @@
 		case "initialized": // notif
 			var params InitializedParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.Initialized(ctx, &params); err != nil {
@@ -89,7 +89,7 @@
 		case "workspace/didChangeConfiguration": // notif
 			var params DidChangeConfigurationParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidChangeConfiguration(ctx, &params); err != nil {
@@ -98,7 +98,7 @@
 		case "textDocument/didOpen": // notif
 			var params DidOpenTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidOpen(ctx, &params); err != nil {
@@ -107,7 +107,7 @@
 		case "textDocument/didChange": // notif
 			var params DidChangeTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidChange(ctx, &params); err != nil {
@@ -116,7 +116,7 @@
 		case "textDocument/didClose": // notif
 			var params DidCloseTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidClose(ctx, &params); err != nil {
@@ -125,7 +125,7 @@
 		case "textDocument/didSave": // notif
 			var params DidSaveTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidSave(ctx, &params); err != nil {
@@ -134,7 +134,7 @@
 		case "textDocument/willSave": // notif
 			var params WillSaveTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.WillSave(ctx, &params); err != nil {
@@ -143,7 +143,7 @@
 		case "workspace/didChangeWatchedFiles": // notif
 			var params DidChangeWatchedFilesParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.DidChangeWatchedFiles(ctx, &params); err != nil {
@@ -152,7 +152,7 @@
 		case "$/setTraceNotification": // notif
 			var params SetTraceParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.SetTraceNotification(ctx, &params); err != nil {
@@ -161,7 +161,7 @@
 		case "$/logTraceNotification": // notif
 			var params LogTraceParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			if err := server.LogTraceNotification(ctx, &params); err != nil {
@@ -170,296 +170,296 @@
 		case "textDocument/implementation": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Implementation(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/typeDefinition": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.TypeDefinition(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/documentColor": // req
 			var params DocumentColorParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.DocumentColor(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/colorPresentation": // req
 			var params ColorPresentationParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.ColorPresentation(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/foldingRange": // req
 			var params FoldingRangeParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.FoldingRange(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/declaration": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Declaration(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "initialize": // req
 			var params InitializeParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Initialize(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "shutdown": // req
 			if r.Params != nil {
-				conn.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
+				r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidParams, "Expected no params"))
 				return
 			}
 			err := server.Shutdown(ctx)
-			if err := conn.Reply(ctx, r, nil, err); err != nil {
+			if err := r.Reply(ctx, nil, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/willSaveWaitUntil": // req
 			var params WillSaveTextDocumentParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.WillSaveWaitUntil(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/completion": // req
 			var params CompletionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Completion(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "completionItem/resolve": // req
 			var params CompletionItem
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Resolve(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/hover": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Hover(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/signatureHelp": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.SignatureHelp(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/definition": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Definition(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/references": // req
 			var params ReferenceParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.References(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/documentHighlight": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.DocumentHighlight(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/documentSymbol": // req
 			var params DocumentSymbolParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.DocumentSymbol(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "workspace/symbol": // req
 			var params WorkspaceSymbolParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Symbol(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/codeAction": // req
 			var params CodeActionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.CodeAction(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/codeLens": // req
 			var params CodeLensParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.CodeLens(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "codeLens/resolve": // req
 			var params CodeLens
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.ResolveCodeLens(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/formatting": // req
 			var params DocumentFormattingParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Formatting(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/rangeFormatting": // req
 			var params DocumentRangeFormattingParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.RangeFormatting(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/onTypeFormatting": // req
 			var params DocumentOnTypeFormattingParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.OnTypeFormatting(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/rename": // req
 			var params RenameParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.Rename(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/prepareRename": // req
 			var params TextDocumentPositionParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.PrepareRename(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "textDocument/documentLink": // req
 			var params DocumentLinkParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.DocumentLink(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "documentLink/resolve": // req
 			var params DocumentLink
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.ResolveDocumentLink(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 		case "workspace/executeCommand": // req
 			var params ExecuteCommandParams
 			if err := json.Unmarshal(*r.Params, &params); err != nil {
-				sendParseError(ctx, log, conn, r, err)
+				sendParseError(ctx, log, r, err)
 				return
 			}
 			resp, err := server.ExecuteCommand(ctx, &params)
-			if err := conn.Reply(ctx, r, resp, err); err != nil {
+			if err := r.Reply(ctx, resp, err); err != nil {
 				log.Errorf(ctx, "%v", err)
 			}
 
 		default:
 			if r.IsNotify() {
-				conn.Reply(ctx, r, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
+				r.Reply(ctx, nil, jsonrpc2.NewErrorf(jsonrpc2.CodeMethodNotFound, "method %q not found", r.Method))
 			}
 		}
 	}