LSP: Provide output.json option for non-VS Code clients (#972)

And use a reader instead of a string for input as suggested by @charlieegan3

Signed-off-by: Anders Eknert <anders@styra.com>

Author: anderseknert

.gitignore

@@ -4,3 +4,9 @@ dist/
 
 /regal
 /regal.exe
+
+# These two files are used by the Regal evaluation Code Lens, where input.json
+# defines the input to use for evaluation, and output.json is where the output
+# ends up unless the client supports presenting it in a different way.
+input.json
+output.json

internal/lsp/eval.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"path"
 	"path/filepath"
@@ -19,7 +20,7 @@ import (
 	"github.com/styrainc/regal/pkg/builtins"
 )
 
-func (l *LanguageServer) Eval(ctx context.Context, query string, input string) (rego.ResultSet, error) {
+func (l *LanguageServer) Eval(ctx context.Context, query string, input io.Reader) (rego.ResultSet, error) {
 	modules := l.cache.GetAllModules()
 	moduleFiles := make([]bundle.ModuleFile, 0, len(modules))
 
@@ -47,10 +48,15 @@ func (l *LanguageServer) Eval(ctx context.Context, query string, input string) (
 		return nil, fmt.Errorf("failed preparing query: %w", err)
 	}
 
-	if input != "" {
+	if input != nil {
 		inputMap := make(map[string]any)
 
-		err = json.Unmarshal([]byte(input), &inputMap)
+		in, err := io.ReadAll(input)
+		if err != nil {
+			return nil, fmt.Errorf("failed reading input: %w", err)
+		}
+
+		err = json.Unmarshal(in, &inputMap)
 		if err != nil {
 			return nil, fmt.Errorf("failed unmarshalling input: %w", err)
 		}
@@ -66,22 +72,23 @@ type EvalPathResult struct {
 	IsUndefined bool `json:"isUndefined"`
 }
 
-func FindInput(file string, workspacePath string) string {
+func FindInput(file string, workspacePath string) io.Reader {
 	relative := strings.TrimPrefix(file, workspacePath)
 	components := strings.Split(path.Dir(relative), string(filepath.Separator))
 
 	for i := range len(components) {
 		inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json")
 
-		if input, err := os.ReadFile(inputPath); err == nil {
-			return string(input)
+		f, err := os.Open(inputPath)
+		if err == nil {
+			return f
 		}
 	}
 
-	return ""
+	return nil
 }
 
-func (l *LanguageServer) EvalWorkspacePath(ctx context.Context, query string, input string) (EvalPathResult, error) {
+func (l *LanguageServer) EvalWorkspacePath(ctx context.Context, query string, input io.Reader) (EvalPathResult, error) {
 	resultQuery := "result := " + query
 
 	result, err := l.Eval(ctx, resultQuery, input)

internal/lsp/eval_test.go

@@ -2,7 +2,9 @@ package lsp
 
 import (
 	"context"
+	"io"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/styrainc/regal/internal/parse"
@@ -39,7 +41,9 @@ func TestEvalWorkspacePath(t *testing.T) {
 	ls.cache.SetModule("file://policy1.rego", module1)
 	ls.cache.SetModule("file://policy2.rego", module2)
 
-	res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", `{"exists": true}`)
+	input := strings.NewReader(`{"exists": true}`)
+
+	res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", input)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -67,15 +71,15 @@ func TestFindInput(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if FindInput(file, workspacePath) != "" {
+	if readInputString(t, file, workspacePath) != "" {
 		t.Fatalf("did not expect to find input.json")
 	}
 
 	content := `{"x": 1}`
 
 	createWithContent(t, tmpDir+"/workspace/foo/bar/input.json", content)
 
-	if res := FindInput(file, workspacePath); res != content {
+	if res := readInputString(t, file, workspacePath); res != content {
 		t.Errorf("expected input at %s, got %s", content, res)
 	}
 
@@ -86,7 +90,7 @@ func TestFindInput(t *testing.T) {
 
 	createWithContent(t, tmpDir+"/workspace/input.json", content)
 
-	if res := FindInput(file, workspacePath); res != content {
+	if res := readInputString(t, file, workspacePath); res != content {
 		t.Errorf("expected input at %s, got %s", content, res)
 	}
 }
@@ -106,3 +110,24 @@ func createWithContent(t *testing.T, path string, content string) {
 		t.Fatal(err)
 	}
 }
+
+func readInputString(t *testing.T, file, workspacePath string) string {
+	t.Helper()
+
+	input := FindInput(file, workspacePath)
+
+	if input == nil {
+		return ""
+	}
+
+	bs, err := io.ReadAll(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if bs == nil {
+		return ""
+	}
+
+	return string(bs)
+}

internal/lsp/server.go

@@ -480,10 +480,8 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) {
 					break
 				}
 
-				input := FindInput(
-					uri.ToPath(l.clientIdentifier, file),
-					uri.ToPath(l.clientIdentifier, l.workspaceRootURI),
-				)
+				workspacePath := uri.ToPath(l.clientIdentifier, l.workspaceRootURI)
+				input := FindInput(uri.ToPath(l.clientIdentifier, file), workspacePath)
 
 				result, err := l.EvalWorkspacePath(ctx, path, input)
 				if err != nil {
@@ -492,16 +490,43 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) {
 					break
 				}
 
-				responseParams := map[string]any{
-					"result": result,
-					"line":   line,
-				}
+				if l.clientIdentifier == clients.IdentifierVSCode {
+					responseParams := map[string]any{
+						"result": result,
+						"line":   line,
+					}
 
-				responseResult := map[string]any{}
+					responseResult := map[string]any{}
 
-				err = l.conn.Call(ctx, "regal/showEvalResult", responseParams, &responseResult)
-				if err != nil {
-					l.logError(fmt.Errorf("failed %s notify: %v", "regal/hello", err.Error()))
+					err = l.conn.Call(ctx, "regal/showEvalResult", responseParams, &responseResult)
+					if err != nil {
+						l.logError(fmt.Errorf("failed %s notify: %v", "regal/hello", err.Error()))
+					}
+				} else {
+					output := filepath.Join(workspacePath, "output.json")
+
+					var f *os.File
+
+					f, err = os.OpenFile(output, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o755)
+
+					if err == nil {
+						var jsonVal []byte
+
+						value := result.Value
+						if result.IsUndefined {
+							// Display undefined as an empty object
+							// we could also go with "<undefined>" or similar
+							value = make(map[string]any)
+						}
+
+						jsonVal, err = json.MarshalIndent(value, "", "  ")
+						if err == nil {
+							// staticcheck thinks err here is never used, but I think that's false?
+							_, err = f.Write(jsonVal) //nolint:staticcheck
+						}
+
+						f.Close()
+					}
 				}
 			}
 
@@ -1000,13 +1025,6 @@ func (l *LanguageServer) handleTextDocumentCodeLens(
 		return nil, fmt.Errorf("failed to unmarshal params: %w", err)
 	}
 
-	if l.clientIdentifier != clients.IdentifierVSCode {
-		// only VSCode has the client side capability to handle the callback request
-		// to handle the result of evaluation, so displaying code lenses for any other
-		// editor is likely just going to result in a bad experience
-		return nil, nil // return a null response, as per the spec
-	}
-
 	module, ok := l.cache.GetModule(params.TextDocument.URI)
 	if !ok {
 		l.logError(fmt.Errorf("failed to get module for uri %q", params.TextDocument.URI))
@@ -1736,16 +1754,10 @@ func (l *LanguageServer) handleInitialize(
 					LabelDetailsSupport: true,
 				},
 			},
+			CodeLensProvider: &types.CodeLensOptions{},
 		},
 	}
 
-	// Since evaluation requires some client side handling, this can't be supported
-	// purely by the LSP. Clients that are capable of handling the code lens callback
-	// should be added here though.
-	if l.clientIdentifier == clients.IdentifierVSCode {
-		initializeResult.Capabilities.CodeLensProvider = &types.CodeLensOptions{}
-	}
-
 	if l.workspaceRootURI != "" {
 		configFile, err := config.FindConfig(uri.ToPath(l.clientIdentifier, l.workspaceRootURI))
 		if err == nil {