Make client requests type safe, unmarshal (#2099)

Author: jakebailey

internal/fourslash/fourslash.go

@@ -353,8 +353,7 @@ func getCapabilitiesWithDefaults(capabilities *lsproto.ClientCapabilities) *lspr
 
 func sendRequest[Params, Resp any](t *testing.T, f *FourslashTest, info lsproto.RequestInfo[Params, Resp], params Params) (*lsproto.Message, Resp, bool) {
 	id := f.nextID()
-	req := lsproto.NewRequestMessage(
-		info.Method,
+	req := info.NewRequestMessage(
 		lsproto.NewID(lsproto.IntegerOrString{Integer: &id}),
 		params,
 	)
@@ -396,8 +395,7 @@ func sendRequest[Params, Resp any](t *testing.T, f *FourslashTest, info lsproto.
 }
 
 func sendNotification[Params any](t *testing.T, f *FourslashTest, info lsproto.NotificationInfo[Params], params Params) {
-	notification := lsproto.NewNotificationMessage(
-		info.Method,
+	notification := info.NewNotificationMessage(
 		params,
 	)
 	f.writeMsg(t, notification.Message())

internal/lsp/lsproto/_generate/generate.mts

@@ -784,6 +784,36 @@ function generateCode() {
     writeLine("}");
     writeLine("");
 
+    // Generate unmarshalResult function
+    writeLine("func unmarshalResult(method Method, data []byte) (any, error) {");
+    writeLine("\tswitch method {");
+
+    // Only requests have results, not notifications
+    for (const request of model.requests) {
+        const methodName = methodNameIdentifier(request.method);
+
+        if (!("result" in request)) {
+            continue;
+        }
+
+        let responseTypeName: string;
+        if (request.typeName && request.typeName.endsWith("Request")) {
+            responseTypeName = request.typeName.replace(/Request$/, "Response");
+        }
+        else {
+            responseTypeName = `${methodName}Response`;
+        }
+
+        writeLine(`\tcase Method${methodName}:`);
+        writeLine(`\t\treturn unmarshalValue[${responseTypeName}](data)`);
+    }
+
+    writeLine("\tdefault:");
+    writeLine(`\t\treturn unmarshalAny(data)`);
+    writeLine("\t}");
+    writeLine("}");
+    writeLine("");
+
     writeLine("// Methods");
     writeLine("const (");
     for (const request of requestsAndNotifications) {

internal/lsp/lsproto/jsonrpc.go

@@ -105,19 +105,20 @@ func (m *Message) UnmarshalJSON(data []byte) error {
 		Method  Method         `json:"method"`
 		ID      *ID            `json:"id,omitzero"`
 		Params  jsontext.Value `json:"params"`
-		Result  any            `json:"result,omitzero"`
-		Error   *ResponseError `json:"error,omitzero"`
+		// We don't have a method in the response, so we have no idea what to decode.
+		// Store the raw text and let the caller decode it.
+		Result jsontext.Value `json:"result,omitzero"`
+		Error  *ResponseError `json:"error,omitzero"`
 	}
 	if err := json.Unmarshal(data, &raw); err != nil {
 		return fmt.Errorf("%w: %w", ErrInvalidRequest, err)
 	}
 	if raw.ID != nil && raw.Method == "" {
 		m.Kind = MessageKindResponse
 		m.msg = &ResponseMessage{
-			JSONRPC: raw.JSONRPC,
-			ID:      raw.ID,
-			Result:  raw.Result,
-			Error:   raw.Error,
+			ID:     raw.ID,
+			Result: raw.Result,
+			Error:  raw.Error,
 		}
 		return nil
 	}
@@ -138,10 +139,9 @@ func (m *Message) UnmarshalJSON(data []byte) error {
 	}
 
 	m.msg = &RequestMessage{
-		JSONRPC: raw.JSONRPC,
-		ID:      raw.ID,
-		Method:  raw.Method,
-		Params:  params,
+		ID:     raw.ID,
+		Method: raw.Method,
+		Params: params,
 	}
 
 	return nil
@@ -151,29 +151,13 @@ func (m *Message) MarshalJSON() ([]byte, error) {
 	return json.Marshal(m.msg)
 }
 
-func NewNotificationMessage(method Method, params any) *RequestMessage {
-	return &RequestMessage{
-		JSONRPC: JSONRPCVersion{},
-		Method:  method,
-		Params:  params,
-	}
-}
-
 type RequestMessage struct {
 	JSONRPC JSONRPCVersion `json:"jsonrpc"`
 	ID      *ID            `json:"id,omitzero"`
 	Method  Method         `json:"method"`
 	Params  any            `json:"params,omitzero"`
 }
 
-func NewRequestMessage(method Method, id *ID, params any) *RequestMessage {
-	return &RequestMessage{
-		ID:     id,
-		Method: method,
-		Params: params,
-	}
-}
-
 func (r *RequestMessage) Message() *Message {
 	return &Message{
 		Kind: MessageKindRequest,

internal/lsp/lsproto/lsp.go

@@ -71,6 +71,14 @@ func unmarshalPtrTo[T any](data []byte) (*T, error) {
 	return &v, nil
 }
 
+func unmarshalValue[T any](data []byte) (T, error) {
+	var v T
+	if err := json.Unmarshal(data, &v); err != nil {
+		return *new(T), fmt.Errorf("failed to unmarshal %T: %w", (*T)(nil), err)
+	}
+	return v, nil
+}
+
 func unmarshalAny(data []byte) (any, error) {
 	var v any
 	if err := json.Unmarshal(data, &v); err != nil {
@@ -129,11 +137,43 @@ type RequestInfo[Params, Resp any] struct {
 	Method Method
 }
 
+func (info RequestInfo[Params, Resp]) UnmarshalResult(result any) (Resp, error) {
+	if r, ok := result.(Resp); ok {
+		return r, nil
+	}
+
+	raw, ok := result.(jsontext.Value)
+	if !ok {
+		return *new(Resp), fmt.Errorf("expected jsontext.Value, got %T", result)
+	}
+
+	r, err := unmarshalResult(info.Method, raw)
+	if err != nil {
+		return *new(Resp), err
+	}
+	return r.(Resp), nil
+}
+
+func (info RequestInfo[Params, Resp]) NewRequestMessage(id *ID, params Params) *RequestMessage {
+	return &RequestMessage{
+		ID:     id,
+		Method: info.Method,
+		Params: params,
+	}
+}
+
 type NotificationInfo[Params any] struct {
 	_      [0]Params
 	Method Method
 }
 
+func (info NotificationInfo[Params]) NewNotificationMessage(params Params) *RequestMessage {
+	return &RequestMessage{
+		Method: info.Method,
+		Params: params,
+	}
+}
+
 type Null struct{}
 
 func (Null) UnmarshalJSONFrom(dec *jsontext.Decoder) error {