fix: add CustomSchemaProvider

internal/adapter/yamlls/client.go

@@ -2,6 +2,7 @@ package yamlls
 
 import (
 	"context"
+	"fmt"
 
 	"go.lsp.dev/protocol"
 )
@@ -13,6 +14,7 @@ func (y Connector) ApplyEdit(ctx context.Context, params *protocol.ApplyWorkspac
 
 // LogMessage implements protocol.Client.
 func (y Connector) LogMessage(ctx context.Context, params *protocol.LogMessageParams) (err error) {
+	logger.Debug(fmt.Sprintf("LogMessage from yamlls: %s %s", params.Type, params.Message))
 	return nil
 }
 

internal/adapter/yamlls/custom_new_client.go

@@ -2,9 +2,11 @@ package yamlls
 
 import (
 	"context"
+	"encoding/json"
 
 	"go.lsp.dev/jsonrpc2"
 	"go.lsp.dev/protocol"
+	"go.lsp.dev/uri"
 	"go.uber.org/zap"
 )
 
@@ -19,6 +21,53 @@ var DefaultCustomHandler = CustomHandler{
 	jsonrpc2.MethodNotFoundHandler, func(ctx context.Context, conn jsonrpc2.Conn) error { return nil },
 }
 
+func NewCustomSchemaHandler(handler jsonrpc2.Handler) *CustomHandler {
+	customHandler := &CustomHandler{
+		Handler: handler,
+		PostInitialize: func(ctx context.Context, conn jsonrpc2.Conn) error {
+			return conn.Notify(ctx, "yaml/registerCustomSchemaRequest", nil)
+		},
+	}
+
+	return customHandler
+}
+
+type CustomSchemaProvider func(ctx context.Context, uri uri.URI) (uri.URI, error)
+
+func NewCustomSchemaProviderHandler(provider CustomSchemaProvider) jsonrpc2.Handler {
+	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
+		switch req.Method() {
+		case "custom/schema/request":
+
+			params := []string{}
+			jsonBytes, err := req.Params().MarshalJSON()
+			if err != nil {
+				logger.Error(err)
+				return reply(ctx, nil, nil)
+			}
+
+			err = json.Unmarshal(jsonBytes, &params)
+			if err != nil {
+				logger.Error(err)
+				return reply(ctx, nil, nil)
+			}
+
+			logger.Println("YamlHandler: custom/schema/request", req.Params())
+
+			if len(params) == 0 {
+				return reply(ctx, nil, nil)
+			}
+
+			schemaURI, err := provider(ctx, uri.New(params[0]))
+			if err != nil {
+				return reply(ctx, nil, err)
+			}
+			return reply(ctx, schemaURI, nil)
+		}
+		return jsonrpc2.MethodNotFoundHandler(ctx, reply, req)
+	}
+}
+
 // CustomNewClient returns the context in which Client is embedded, jsonrpc2.Conn, and the Server.
 func (yamllsConnector Connector) CustomNewClient(ctx context.Context, client protocol.Client, stream jsonrpc2.Stream, logger *zap.Logger) (context.Context, jsonrpc2.Conn, protocol.Server) {
 	ctx = protocol.WithClient(ctx, client)

internal/handler/yaml_handler/configure.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/mrjosh/helm-ls/internal/adapter/yamlls"
 	"github.com/mrjosh/helm-ls/internal/util"
-	"go.lsp.dev/jsonrpc2"
 )
 
 func (h *YamlHandler) Configure(ctx context.Context, helmlsConfig util.HelmlsConfiguration) {
@@ -15,18 +14,15 @@ func (h *YamlHandler) Configure(ctx context.Context, helmlsConfig util.HelmlsCon
 func (h *YamlHandler) configureYamlls(ctx context.Context, helmlsConfig util.HelmlsConfiguration) {
 	config := helmlsConfig
 
-	customHandler := yamlls.CustomHandler{
-		Handler: h.CustomHandler,
-		PostInitialize: func(ctx context.Context, conn jsonrpc2.Conn) error {
-			return conn.Notify(ctx, "yaml/registerCustomSchemaRequest", nil)
-		},
-	}
-
 	connector := yamlls.NewConnector(ctx,
 		config.YamllsConfiguration,
 		h.client,
 		h.documents,
-		customHandler)
+		*yamlls.NewCustomSchemaHandler(
+			yamlls.NewCustomSchemaProviderHandler(h.CustomSchemaProvider),
+		),
+	)
+
 	h.setYamllsConnector(connector)
 
 	err := h.yamllsConnector.CallInitialize(ctx, h.chartStore.RootURI)

internal/handler/yaml_handler/yaml_handler.go

@@ -52,6 +52,19 @@ func (h *YamlHandler) setYamllsConnector(yamllsConnector *yamlls.Connector) {
 	h.yamllsConnector = yamllsConnector
 }
 
+func (h *YamlHandler) CustomSchemaProvider(ctx context.Context, URI uri.URI) (uri.URI, error) {
+	chart, err := h.chartStore.GetChartForDoc(URI)
+	if err != nil {
+		logger.Error(err)
+	}
+	schemaFilePath, err := jsonschema.CreateJsonSchemaForChart(chart)
+	if err != nil {
+		logger.Error(err)
+		return uri.New(""), err
+	}
+	return uri.File(schemaFilePath), nil
+}
+
 func (h *YamlHandler) CustomHandler(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) error {
 	switch req.Method() {
 	case "custom/schema/request":