refactor: move autofix and autofix feedback logic from language server to code-client-go [IDE-727] (#100)

Author: andrewrobinsonhodges-snyk

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/go-git/go-git/v5 v5.14.0
 	github.com/golang/mock v1.6.0
 	github.com/google/uuid v1.6.0
+	github.com/hexops/gotextdiff v1.0.3
 	github.com/oapi-codegen/oapi-codegen/v2 v2.4.1
 	github.com/oapi-codegen/runtime v1.1.1
 	github.com/pact-foundation/pact-go/v2 v2.0.5

go.sum

@@ -82,6 +82,8 @@ github.com/hashicorp/go-version v1.6.0 h1:feTTfFNnjP967rlCxM/I9g701jU+RN74YKx2mO
 github.com/hashicorp/go-version v1.6.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
 github.com/hashicorp/logutils v1.0.0 h1:dLEQVugN8vlakKOUE3ihGLTZJRB4j+M2cdTm/ORI65Y=
 github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
+github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
+github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=

llm/api_client.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"net/url"
@@ -40,18 +41,42 @@ func (d *DeepCodeLLMBindingImpl) runExplain(ctx context.Context, options Explain
 			return Explanations{}, err
 		}
 	}
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), bytes.NewBuffer(requestBody))
+
+	responseBody, err := d.submitRequest(ctx, u, requestBody)
 	if err != nil {
-		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error creating request")
 		return Explanations{}, err
 	}
 
+	var response explainResponse
+	var explains Explanations
+	response.Status = completeStatus
+	err = json.Unmarshal(responseBody, &response)
+	if err != nil {
+		logger.Err(err).Str("responseBody", string(responseBody)).Msg("error unmarshalling")
+		return Explanations{}, err
+	}
+
+	explains = response.Explanation
+
+	return explains, nil
+}
+
+func (d *DeepCodeLLMBindingImpl) submitRequest(ctx context.Context, url *url.URL, requestBody []byte) ([]byte, error) {
+	logger := d.logger.With().Str("method", "submitRequest").Logger()
+	logger.Trace().Str("payload body: %s\n", string(requestBody)).Msg("Marshaled payload")
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url.String(), bytes.NewBuffer(requestBody))
+	if err != nil {
+		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error creating request")
+		return nil, err
+	}
+
 	d.addDefaultHeaders(req)
 
 	resp, err := d.httpClientFunc().Do(req) //nolint:bodyclose // this seems to be a false positive
 	if err != nil {
 		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error getting response")
-		return Explanations{}, err
+		return nil, err
 	}
 	defer func(Body io.ReadCloser) {
 		bodyCloseErr := Body.Close()
@@ -64,21 +89,11 @@ func (d *DeepCodeLLMBindingImpl) runExplain(ctx context.Context, options Explain
 	responseBody, err := io.ReadAll(resp.Body)
 	if err != nil {
 		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error reading all response")
-		return Explanations{}, err
+		return nil, err
 	}
 	logger.Debug().Str("response body: %s\n", string(responseBody)).Msg("Got the response")
-	var response explainResponse
-	var explains Explanations
-	response.Status = completeStatus
-	err = json.Unmarshal(responseBody, &response)
-	if err != nil {
-		logger.Err(err).Str("responseBody", string(responseBody)).Msg("error unmarshalling")
-		return Explanations{}, err
-	}
-
-	explains = response.Explanation
 
-	return explains, nil
+	return responseBody, nil
 }
 
 func (d *DeepCodeLLMBindingImpl) explainRequestBody(options *ExplainOptions) ([]byte, error) {
@@ -105,6 +120,111 @@ func (d *DeepCodeLLMBindingImpl) explainRequestBody(options *ExplainOptions) ([]
 	return requestBody, marshalErr
 }
 
+var failed = AutofixStatus{Message: "FAILED"}
+
+func (d *DeepCodeLLMBindingImpl) runAutofix(ctx context.Context, requestId string, options AutofixOptions) (AutofixResponse, AutofixStatus, error) {
+	span := d.instrumentor.StartSpan(ctx, "code.RunAutofix")
+	defer span.Finish()
+
+	logger := d.logger.With().Str("method", "code.RunAutofix").Str("requestId", requestId).Logger()
+
+	requestBody, err := d.autofixRequestBody(&options)
+	if err != nil {
+		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error creating request body")
+		return AutofixResponse{}, failed, err
+	}
+
+	logger.Info().Msg("Started obtaining autofix Response")
+	responseBody, err := d.submitRequest(ctx, options.Endpoint, requestBody)
+	logger.Info().Msg("Finished obtaining autofix Response")
+
+	if err != nil {
+		logger.Err(err).Str("responseBody", string(responseBody)).Msg("error response from autofix")
+		return AutofixResponse{}, failed, err
+	}
+
+	var response AutofixResponse
+	err = json.Unmarshal(responseBody, &response)
+	if err != nil {
+		logger.Err(err).Str("responseBody", string(responseBody)).Msg("error unmarshalling")
+		return AutofixResponse{}, failed, err
+	}
+
+	logger.Debug().Msgf("Status: %s", response.Status)
+
+	if response.Status == failed.Message {
+		errMsg := "autofix failed"
+		logger.Error().Str("responseStatus", response.Status).Msg(errMsg)
+		return response, failed, errors.New(errMsg)
+	}
+
+	if response.Status == "" {
+		errMsg := "unknown response status (empty)"
+		logger.Error().Str("responseStatus", response.Status).Msg(errMsg)
+		return response, failed, errors.New(errMsg)
+	}
+
+	status := AutofixStatus{Message: response.Status}
+	if response.Status != completeStatus {
+		return response, status, nil
+	}
+
+	return response, status, nil
+}
+
+func (d *DeepCodeLLMBindingImpl) autofixRequestBody(options *AutofixOptions) ([]byte, error) {
+	request := AutofixRequest{
+		Key: AutofixRequestKey{
+			Type:     "file",
+			Hash:     options.BundleHash,
+			FilePath: options.FilePath,
+			RuleId:   options.RuleID,
+			LineNum:  options.LineNum,
+		},
+		AnalysisContext:     options.CodeRequestContext,
+		IdeExtensionDetails: options.IdeExtensionDetails,
+	}
+	if len(options.ShardKey) > 0 {
+		request.Key.Shard = options.ShardKey
+	}
+
+	requestBody, err := json.Marshal(request)
+	return requestBody, err
+}
+
+func (d *DeepCodeLLMBindingImpl) submitAutofixFeedback(ctx context.Context, requestId string, options AutofixFeedbackOptions) error {
+	span := d.instrumentor.StartSpan(ctx, "code.SubmitAutofixFeedback")
+	defer span.Finish()
+
+	logger := d.logger.With().Str("method", "code.SubmitAutofixFeedback").Str("requestId", requestId).Logger()
+
+	requestBody, err := d.autofixFeedbackRequestBody(&options)
+	if err != nil {
+		logger.Err(err).Str("requestBody", string(requestBody)).Msg("error creating request body")
+		return err
+	}
+
+	logger.Info().Msg("Started obtaining autofix Response")
+	_, err = d.submitRequest(ctx, options.Endpoint, requestBody)
+	logger.Info().Msg("Finished obtaining autofix Response")
+
+	return err
+}
+
+func (d *DeepCodeLLMBindingImpl) autofixFeedbackRequestBody(options *AutofixFeedbackOptions) ([]byte, error) {
+	request := AutofixUserEvent{
+		Channel:             "IDE",
+		EventType:           options.Result,
+		EventDetails:        AutofixEventDetails{FixId: options.FixID},
+		AnalysisContext:     options.CodeRequestContext,
+		IdeExtensionDetails: options.IdeExtensionDetails,
+	}
+
+	requestBody, err := json.Marshal(request)
+
+	return requestBody, err
+}
+
 func prepareDiffs(diffs []string) []string {
 	cleanedDiffs := make([]string, 0, len(diffs))
 	for _, diff := range diffs {

llm/binding.go

@@ -51,17 +51,45 @@ type ExplainResult []string
 type DeepCodeLLMBinding interface {
 	SnykLLMBindings
 	ExplainWithOptions(ctx context.Context, options ExplainOptions) (ExplainResult, error)
+	GetAutofixDiffs(ctx context.Context, baseDir string, options AutofixOptions) (unifiedDiffSuggestions []AutofixUnifiedDiffSuggestion, status AutofixStatus, err error)
+	SubmitAutofixFeedback(ctx context.Context, requestId string, options AutofixFeedbackOptions) error
 }
 
 // DeepCodeLLMBindingImpl is an LLM binding for the Snyk Code LLM.
-// Currently, it only supports explain.
 type DeepCodeLLMBindingImpl struct {
 	httpClientFunc func() http.HTTPClient
 	logger         *zerolog.Logger
 	outputFormat   OutputFormat
 	instrumentor   observability.Instrumentor
 }
 
+func (d *DeepCodeLLMBindingImpl) SubmitAutofixFeedback(ctx context.Context, requestId string, options AutofixFeedbackOptions) error {
+	method := "SubmitAutofixFeedback"
+	span := d.instrumentor.StartSpan(ctx, method)
+	defer d.instrumentor.Finish(span)
+	logger := d.logger.With().Str("method", method).Str("requestId", requestId).Logger()
+	logger.Info().Msg("Started submitting autofix feedback")
+	defer logger.Info().Msg("Finished submitting autofix feedback")
+
+	err := d.submitAutofixFeedback(ctx, requestId, options)
+	return err
+}
+
+func (d *DeepCodeLLMBindingImpl) GetAutofixDiffs(ctx context.Context, requestId string, options AutofixOptions) (unifiedDiffSuggestions []AutofixUnifiedDiffSuggestion, status AutofixStatus, err error) {
+	method := "GetAutofixDiffs"
+	span := d.instrumentor.StartSpan(ctx, method)
+	defer d.instrumentor.Finish(span)
+	logger := d.logger.With().Str("method", method).Str("requestId", requestId).Logger()
+	logger.Info().Msg("Started obtaining autofix diffs")
+	defer logger.Info().Msg("Finished obtaining autofix diffs")
+
+	autofixResponse, status, err := d.runAutofix(ctx, requestId, options)
+	if err != nil {
+		return nil, status, err
+	}
+	return autofixResponse.toUnifiedDiffSuggestions(d.logger, options.BaseDir, options.FilePath), status, err
+}
+
 func (d *DeepCodeLLMBindingImpl) ExplainWithOptions(ctx context.Context, options ExplainOptions) (ExplainResult, error) {
 	s := d.instrumentor.StartSpan(ctx, "code.ExplainWithOptions")
 	defer d.instrumentor.Finish(s)

llm/convert.go

@@ -0,0 +1,57 @@
+package llm
+
+import (
+	"fmt"
+	"net/url"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/hexops/gotextdiff"
+	"github.com/hexops/gotextdiff/myers"
+	"github.com/hexops/gotextdiff/span"
+	"github.com/rs/zerolog"
+)
+
+func (s *AutofixResponse) toUnifiedDiffSuggestions(logger *zerolog.Logger, baseDir string, filePath string) []AutofixUnifiedDiffSuggestion {
+	var fixSuggestions []AutofixUnifiedDiffSuggestion
+	for _, suggestion := range s.AutofixSuggestions {
+		decodedPath, unifiedDiff := getPathAndUnifiedDiff(logger, baseDir, filePath, suggestion.Value)
+		if decodedPath == "" || unifiedDiff == "" {
+			continue
+		}
+
+		d := AutofixUnifiedDiffSuggestion{
+			FixId:               suggestion.Id,
+			UnifiedDiffsPerFile: map[string]string{},
+		}
+
+		d.UnifiedDiffsPerFile[decodedPath] = unifiedDiff
+		fixSuggestions = append(fixSuggestions, d)
+	}
+	return fixSuggestions
+}
+
+func getPathAndUnifiedDiff(zeroLogger *zerolog.Logger, baseDir string, filePath string, newText string) (decodedPath string, unifiedDiff string) {
+	logger := zeroLogger.With().Str("method", "getUnifiedDiff").Logger()
+
+	decodedPath, err := url.PathUnescape(filepath.Join(baseDir, filePath))
+	if err != nil {
+		logger.Err(err).Msgf("cannot decode filePath %s", filePath)
+		return
+	}
+	logger.Debug().Msgf("File decodedPath %s", decodedPath)
+
+	fileContent, err := os.ReadFile(decodedPath)
+	if err != nil {
+		logger.Err(err).Msgf("cannot read fileContent %s", decodedPath)
+		return
+	}
+
+	// Workaround: AI Suggestion API only returns \n new lines. It doesn't consider carriage returns.
+	contentBefore := strings.Replace(string(fileContent), "\r\n", "\n", -1)
+	edits := myers.ComputeEdits(span.URIFromPath(decodedPath), contentBefore, newText)
+	unifiedDiff = fmt.Sprint(gotextdiff.ToUnified(decodedPath, decodedPath+"fixed", contentBefore, edits))
+
+	return decodedPath, unifiedDiff
+}