diff --git a/llm/api_client.go b/llm/api_client.go index 9b6f3dc..465127d 100644 --- a/llm/api_client.go +++ b/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,9 +62,11 @@ 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 { @@ -72,7 +74,7 @@ func (d *DeepCodeLLMBindingImpl) submitRequest(ctx context.Context, url *url.URL 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") } diff --git a/llm/api_client_test.go b/llm/api_client_test.go index 4bf169b..704c8b1 100644 --- a/llm/api_client_test.go +++ b/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") diff --git a/llm/binding.go b/llm/binding.go index 01b7720..cfdc9f4 100644 --- a/llm/binding.go +++ b/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 }