lsp: add action and handlers for support vscode-based opa-explorer (#1862)

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

Author: srenatus

internal/lsp/server.go

@@ -26,6 +26,7 @@ import (
 	"github.com/open-policy-agent/regal/bundle"
 	"github.com/open-policy-agent/regal/internal/capabilities"
 	"github.com/open-policy-agent/regal/internal/compile"
+	"github.com/open-policy-agent/regal/internal/explorer"
 	rio "github.com/open-policy-agent/regal/internal/io"
 	"github.com/open-policy-agent/regal/internal/io/files"
 	"github.com/open-policy-agent/regal/internal/lsp/bundles"
@@ -569,6 +570,35 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) {
 		case <-ctx.Done():
 			return
 		case params := <-l.commandRequest:
+			// Handle regal.explorer separately as VS Code sends map arguments
+			if params.Command == "regal.explorer" {
+				var explorerArgs types.ExplorerCommandArgs
+
+				if len(params.Arguments) > 0 {
+					arg, ok := params.Arguments[0].(map[string]any)
+					if !ok {
+						l.log.Message("expected explorer argument to be a map, got %T", params.Arguments[0])
+
+						continue
+					}
+
+					explorerArgs = types.ExplorerCommandArgs{
+						Target:      util.GetMapValue[string](arg, "target"),
+						Strict:      util.GetMapValue[bool](arg, "strict"),
+						Annotations: util.GetMapValue[bool](arg, "annotations"),
+						Print:       util.GetMapValue[bool](arg, "print"),
+						Format:      util.GetMapValue[bool](arg, "format"),
+					}
+				}
+
+				if err := l.handleExplorerCommand(ctx, explorerArgs); err != nil {
+					l.log.Message("failed to handle explorer command: %s", err)
+				}
+
+				continue
+			}
+
+			// Handle all other commands (they use string arguments)
 			var (
 				editParams *types.ApplyWorkspaceEditParams
 				args       types.CommandArgs
@@ -1919,6 +1949,7 @@ func (l *LanguageServer) handleInitialize(ctx context.Context, params types.Init
 				Commands: []string{
 					"regal.debug",
 					"regal.eval",
+					"regal.explorer",
 					"regal.fix.opa-fmt",
 					"regal.fix.use-rego-v1",
 					"regal.fix.use-assignment-operator",
@@ -2283,6 +2314,68 @@ func (l *LanguageServer) regalContext(fileURI string, _ *rego.Requirements) *reg
 	}
 }
 
+func (l *LanguageServer) handleExplorerCommand(ctx context.Context, args types.ExplorerCommandArgs) error {
+	if args.Target == "" {
+		l.log.Message("expected command target, got empty string")
+
+		return errors.New("target file URI is required")
+	}
+
+	contents, ok := l.cache.GetFileContents(args.Target)
+	if !ok {
+		return fmt.Errorf("could not get file contents for uri %q", args.Target)
+	}
+
+	path := l.toRelativePath(args.Target)
+
+	compileResults := explorer.CompilerStages(
+		path,
+		contents,
+		args.Strict,
+		args.Annotations,
+		args.Print,
+	)
+
+	stages := make([]types.ExplorerStageResult, 0, len(compileResults))
+	hasErrors := false
+
+	for _, cs := range compileResults {
+		stage := types.ExplorerStageResult{
+			Name:  cs.Stage,
+			Error: cs.Error != "",
+		}
+
+		if cs.Error != "" {
+			hasErrors = true
+			stage.Output = cs.Error
+		} else {
+			if args.Format {
+				stage.Output = cs.FormattedResult()
+			} else if cs.Result != nil {
+				stage.Output = cs.Result.String()
+			}
+		}
+
+		stages = append(stages, stage)
+	}
+
+	responseParams := types.ExplorerResult{
+		Stages: stages,
+	}
+
+	if !hasErrors {
+		if plan, err := explorer.Plan(ctx, path, contents, args.Print); err == nil {
+			responseParams.Plan = plan
+		}
+	}
+
+	if err := l.conn.Notify(ctx, "regal/showExplorerResult", responseParams); err != nil {
+		return fmt.Errorf("regal/showExplorerResult notification failed: %w", err)
+	}
+
+	return nil
+}
+
 func (l *LanguageServer) handleEvalCommand(ctx context.Context, args types.CommandArgs) error {
 	if args.Target == "" || args.Query == "" {
 		l.log.Message("expected command target and query, got target %q, query %q", args.Target, args.Query)

internal/lsp/server_command_test.go

@@ -154,3 +154,118 @@ allow if {
 		})
 	}
 }
+
+func TestExecuteCommandExplorer(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(t.Context())
+
+	defer func() {
+		time.Sleep(200 * time.Millisecond)
+		cancel()
+	}()
+
+	receivedNotifications := make(chan map[string]any, defaultBufferedChannelSize)
+
+	createExplorerNotificationTestHandler := func(
+		t *testing.T,
+		receivedNotifications chan map[string]any,
+	) func(_ context.Context, _ *jsonrpc2.Conn, _ *jsonrpc2.Request) (result any, err error) {
+		t.Helper()
+
+		return func(_ context.Context, _ *jsonrpc2.Conn, req *jsonrpc2.Request) (result any, err error) {
+			if req.Method == "regal/showExplorerResult" {
+				if req.Params == nil {
+					t.Fatal("expected notification params to be non-nil")
+				}
+
+				var notificationData map[string]any
+				if err := encoding.JSON().Unmarshal(*req.Params, &notificationData); err != nil {
+					t.Fatalf("failed to unmarshal notification params: %s", err)
+				}
+
+				receivedNotifications <- notificationData
+
+				return nil, nil
+			}
+
+			if req.Method == "workspace/applyEdit" {
+				return map[string]any{"applied": true}, nil
+			}
+
+			return struct{}{}, nil
+		}
+	}
+
+	tempDir := t.TempDir()
+	clientHandler := createExplorerNotificationTestHandler(t, receivedNotifications)
+	ls, connClient := createAndInitServer(t, ctx, tempDir, clientHandler)
+
+	go ls.StartCommandWorker(ctx)
+
+	content := `package test
+
+allow if {
+	1 == 1
+}
+`
+
+	mainRegoURI := uri.FromPath(clients.IdentifierGoTest, filepath.Join(tempDir, "test.rego"))
+	ls.cache.SetFileContents(mainRegoURI, content)
+
+	executeParams := types.ExecuteCommandParams{
+		Command: "regal.explorer",
+		Arguments: []any{
+			map[string]any{
+				"target":      mainRegoURI,
+				"strict":      false,
+				"annotations": false,
+				"print":       false,
+				"format":      true,
+			},
+		},
+	}
+
+	var executeResponse any
+
+	testutil.NoErr(connClient.Call(ctx, "workspace/executeCommand", executeParams, &executeResponse))(t)
+
+	timeout := time.NewTimer(determineTimeout())
+	defer timeout.Stop()
+
+	select {
+	case notification := <-receivedNotifications:
+		if _, ok := notification["stages"]; !ok {
+			t.Fatal("expected notification to contain 'stages' field")
+		}
+
+		stages, ok := notification["stages"].([]any)
+		if !ok {
+			t.Fatalf("expected 'stages' to be a slice, got %T", notification["stages"])
+		}
+
+		if len(stages) == 0 {
+			t.Fatal("expected at least one compilation stage")
+		}
+
+		firstStage, ok := stages[0].(map[string]any)
+		if !ok {
+			t.Fatalf("expected first stage to be a map, got %T", stages[0])
+		}
+
+		if _, ok := firstStage["name"]; !ok {
+			t.Fatal("expected stage to have 'name' field")
+		}
+
+		if _, ok := firstStage["output"]; !ok {
+			t.Fatal("expected stage to have 'output' field")
+		}
+
+		if _, ok := firstStage["error"]; !ok {
+			t.Fatal("expected stage to have 'error' field")
+		}
+
+	case <-timeout.C:
+		t.Fatal("timeout waiting for regal/showExplorerResult notification")
+	}
+}

internal/lsp/types/types.go

@@ -575,5 +575,24 @@ type (
 		Full   bool                 `json:"full,omitempty"`
 	}
 
+	ExplorerCommandArgs struct {
+		Target      string `json:"target"`
+		Strict      bool   `json:"strict,omitempty"`
+		Annotations bool   `json:"annotations,omitempty"`
+		Print       bool   `json:"print,omitempty"`
+		Format      bool   `json:"format,omitempty"`
+	}
+
+	ExplorerStageResult struct {
+		Name   string `json:"name"`
+		Output string `json:"output"`
+		Error  bool   `json:"error"`
+	}
+
+	ExplorerResult struct {
+		Stages []ExplorerStageResult `json:"stages"`
+		Plan   string                `json:"plan,omitempty"`
+	}
+
 	iuint interface{ ~int | ~uint }
 )

internal/util/util.go

@@ -270,6 +270,20 @@ func Pointer[T any](v T) *T {
 	return &v
 }
 
+// GetMapValue extracts a typed value from a map[string]any, returning the value if the type matched,
+// or the zero values of correct type.
+func GetMapValue[T any](m map[string]any, key string) T {
+	if val, ok := m[key]; ok {
+		if typed, ok := val.(T); ok {
+			return typed
+		}
+	}
+
+	var zero T
+
+	return zero
+}
+
 // AnySliceTo converts a slice of any to a slice of T, returning an error if any element cannot be casted.
 func AnySliceTo[T any](in []any) ([]T, error) {
 	out := make([]T, 0, len(in))