fix: set org header for AI Fix (#106)

Author: ShawkyZ

llm/api_client.go

@@ -43,7 +43,7 @@ func (d *DeepCodeLLMBindingImpl) runExplain(ctx context.Context, options Explain
 		}
 	}
 
-	responseBody, err := d.submitRequest(ctx, u, requestBody)
+	responseBody, err := d.submitRequest(span.Context(), u, requestBody, "")
 	if err != nil {
 		return Explanations{}, err
 	}
@@ -62,17 +62,19 @@ func (d *DeepCodeLLMBindingImpl) runExplain(ctx context.Context, options Explain
 	return explains, nil
 }
 
-func (d *DeepCodeLLMBindingImpl) submitRequest(ctx context.Context, url *url.URL, requestBody []byte) ([]byte, error) {
+func (d *DeepCodeLLMBindingImpl) submitRequest(ctx context.Context, url *url.URL, requestBody []byte, orgId string) ([]byte, error) {
 	logger := d.logger.With().Str("method", "submitRequest").Logger()
 	logger.Trace().Str("payload body: %s\n", string(requestBody)).Msg("Marshaled payload")
+	span := d.instrumentor.StartSpan(ctx, "code.SubmitRequest")
+	defer span.Finish()
 
 	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)
+	d.addDefaultHeaders(req, span.GetTraceId(), orgId)
 
 	resp, err := d.httpClientFunc().Do(req) //nolint:bodyclose // this seems to be a false positive
 	if err != nil {
@@ -123,11 +125,11 @@ func (d *DeepCodeLLMBindingImpl) explainRequestBody(options *ExplainOptions) ([]
 
 var failed = AutofixStatus{Message: "FAILED"}
 
-func (d *DeepCodeLLMBindingImpl) runAutofix(ctx context.Context, requestId string, options AutofixOptions) (AutofixResponse, AutofixStatus, error) {
+func (d *DeepCodeLLMBindingImpl) runAutofix(ctx context.Context, 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()
+	logger := d.logger.With().Str("method", "code.RunAutofix").Str("requestId", span.GetTraceId()).Logger()
 
 	endpoint, err := url.Parse(fmt.Sprintf("%s/autofix/suggestions", options.Host))
 	if err != nil {
@@ -142,7 +144,7 @@ func (d *DeepCodeLLMBindingImpl) runAutofix(ctx context.Context, requestId strin
 	}
 
 	logger.Info().Msg("Started obtaining autofix Response")
-	responseBody, err := d.submitRequest(ctx, endpoint, requestBody)
+	responseBody, err := d.submitRequest(span.Context(), endpoint, requestBody, options.CodeRequestContext.Org.PublicId)
 	logger.Info().Msg("Finished obtaining autofix Response")
 
 	if err != nil {
@@ -199,11 +201,11 @@ func (d *DeepCodeLLMBindingImpl) autofixRequestBody(options *AutofixOptions) ([]
 	return requestBody, err
 }
 
-func (d *DeepCodeLLMBindingImpl) submitAutofixFeedback(ctx context.Context, requestId string, options AutofixFeedbackOptions) error {
+func (d *DeepCodeLLMBindingImpl) submitAutofixFeedback(ctx context.Context, 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()
+	logger := d.logger.With().Str("method", "code.SubmitAutofixFeedback").Str("requestId", span.GetTraceId()).Logger()
 
 	endpoint, err := url.Parse(fmt.Sprintf("%s/autofix/event", options.Host))
 	if err != nil {
@@ -218,7 +220,7 @@ func (d *DeepCodeLLMBindingImpl) submitAutofixFeedback(ctx context.Context, requ
 	}
 
 	logger.Info().Msg("Started obtaining autofix Response")
-	_, err = d.submitRequest(ctx, endpoint, requestBody)
+	_, err = d.submitRequest(span.Context(), endpoint, requestBody, options.CodeRequestContext.Org.PublicId)
 	logger.Info().Msg("Finished obtaining autofix Response")
 
 	return err
@@ -257,7 +259,14 @@ func prepareDiffs(diffs []string) []string {
 	return encodedDiffs
 }
 
-func (d *DeepCodeLLMBindingImpl) addDefaultHeaders(req *http.Request) {
+func (d *DeepCodeLLMBindingImpl) addDefaultHeaders(req *http.Request, requestId string, orgId string) {
+	// if requestId is empty it will be enriched from the Gateway
+	if len(requestId) > 0 {
+		req.Header.Set("snyk-request-id", requestId)
+	}
+	if len(orgId) > 0 {
+		req.Header.Set("snyk-org-name", orgId)
+	}
 	req.Header.Set("Cache-Control", "private, max-age=0, no-cache")
 	req.Header.Set("Content-Type", "application/json")
 }

llm/api_client_test.go

@@ -265,7 +265,7 @@ func TestAddDefaultHeadersWithExistingHeaders(t *testing.T) {
 	d := &DeepCodeLLMBindingImpl{} // Initialize your struct if needed
 	req := &http.Request{Header: http.Header{"Existing-Header": {"existing-value"}}}
 
-	d.addDefaultHeaders(req)
+	d.addDefaultHeaders(req, "", "")
 
 	cacheControl := req.Header.Get("Cache-Control")
 	contentType := req.Header.Get("Content-Type")

llm/binding.go

@@ -63,27 +63,28 @@ type DeepCodeLLMBindingImpl struct {
 	instrumentor   observability.Instrumentor
 }
 
-func (d *DeepCodeLLMBindingImpl) SubmitAutofixFeedback(ctx context.Context, requestId string, options AutofixFeedbackOptions) error {
+func (d *DeepCodeLLMBindingImpl) SubmitAutofixFeedback(ctx context.Context, fixId 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 := d.logger.With().Str("method", method).Str("fixId", fixId).Logger()
 	logger.Info().Msg("Started submitting autofix feedback")
 	defer logger.Info().Msg("Finished submitting autofix feedback")
 
-	err := d.submitAutofixFeedback(ctx, requestId, options)
+	err := d.submitAutofixFeedback(span.Context(), options)
 	return err
 }
 
-func (d *DeepCodeLLMBindingImpl) GetAutofixDiffs(ctx context.Context, requestId string, options AutofixOptions) (unifiedDiffSuggestions []AutofixUnifiedDiffSuggestion, status AutofixStatus, err error) {
+func (d *DeepCodeLLMBindingImpl) GetAutofixDiffs(ctx context.Context, _ string, options AutofixOptions) (unifiedDiffSuggestions []AutofixUnifiedDiffSuggestion, status AutofixStatus, err error) {
 	method := "GetAutofixDiffs"
 	span := d.instrumentor.StartSpan(ctx, method)
 	defer d.instrumentor.Finish(span)
+	requestId := span.GetTraceId()
 	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)
+	autofixResponse, status, err := d.runAutofix(span.Context(), options)
 	if err != nil {
 		return nil, status, err
 	}