Gateway: forward X-Fireworks-Genie header from client (#63460)

Previously, there was no way to enable the "tracing" feature from
Fireworks https://readme.fireworks.ai/docs/enabling-tracing This PR
solves the problem by forwarding the `X-Fireworks-Genie` HTTP header to
Fireworks if this HTTP header is set by the Gateway client.

Fixes CODY-2555

<!-- 💡 To write a useful PR description, make sure that your description
covers:
- WHAT this PR is changing:
    - How was it PREVIOUSLY.
    - How it will be from NOW on.
- WHY this PR is needed.
- CONTEXT, i.e. to which initiative, project or RFC it belongs.

The structure of the description doesn't matter as much as covering
these points, so use
your best judgement based on your context.
Learn how to write good pull request description:
https://www.notion.so/sourcegraph/Write-a-good-pull-request-description-610a7fd3e613496eb76f450db5a49b6e?pvs=4
-->


## Test plan

<!-- All pull requests REQUIRE a test plan:
https://docs-legacy.sourcegraph.com/dev/background-information/testing_principles
-->
N/A

## Changelog

<!--
1. Ensure your pull request title is formatted as: $type($domain): $what
2. Add bullet list items for each additional detail you want to cover
(see example below)
3. You can edit this after the pull request was merged, as long as
release shipping it hasn't been promoted to the public.
4. For more information, please see this how-to
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c?

Audience: TS/CSE > Customers > Teammates (in that order).

Cheat sheet: $type = chore|fix|feat $domain:
source|search|ci|release|plg|cody|local|...
-->

<!--
Example:

Title: fix(search): parse quotes with the appropriate context
Changelog section:

## Changelog

- When a quote is used with regexp pattern type, then ...
- Refactored underlying code.
-->

Author: olafurpg

cmd/cody-gateway/internal/httpapi/completions/anthropic.go

@@ -155,15 +155,15 @@ func (a *AnthropicHandlerMethods) getRequestMetadata(body anthropicRequest) (mod
 	}
 }
 
-func (a *AnthropicHandlerMethods) transformRequest(r *http.Request) {
+func (a *AnthropicHandlerMethods) transformRequest(downstreamRequest, upstreamRequest *http.Request) {
 	// Mimic headers set by the official Anthropic client:
 	// https://sourcegraph.com/github.com/anthropics/anthropic-sdk-typescript@493075d70f50f1568a276ed0cb177e297f5fef9f/-/blob/src/index.ts
-	r.Header.Set("Cache-Control", "no-cache")
-	r.Header.Set("Accept", "application/json")
-	r.Header.Set("Content-Type", "application/json")
-	r.Header.Set("Client", "sourcegraph-cody-gateway/1.0")
-	r.Header.Set("X-API-Key", a.config.AccessToken)
-	r.Header.Set("anthropic-version", "2023-01-01")
+	upstreamRequest.Header.Set("Cache-Control", "no-cache")
+	upstreamRequest.Header.Set("Accept", "application/json")
+	upstreamRequest.Header.Set("Content-Type", "application/json")
+	upstreamRequest.Header.Set("Client", "sourcegraph-cody-gateway/1.0")
+	upstreamRequest.Header.Set("X-API-Key", a.config.AccessToken)
+	upstreamRequest.Header.Set("anthropic-version", "2023-01-01")
 }
 
 func (a *AnthropicHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody anthropicRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {

cmd/cody-gateway/internal/httpapi/completions/anthropicmessages.go

@@ -209,10 +209,10 @@ func (a *AnthropicMessagesHandlerMethods) getRequestMetadata(body anthropicMessa
 	}
 }
 
-func (a *AnthropicMessagesHandlerMethods) transformRequest(r *http.Request) {
-	r.Header.Set("Content-Type", "application/json")
-	r.Header.Set("X-API-Key", a.config.AccessToken)
-	r.Header.Set("anthropic-version", "2023-06-01")
+func (a *AnthropicMessagesHandlerMethods) transformRequest(downstreamRequest, upstreamRequest *http.Request) {
+	upstreamRequest.Header.Set("Content-Type", "application/json")
+	upstreamRequest.Header.Set("X-API-Key", a.config.AccessToken)
+	upstreamRequest.Header.Set("anthropic-version", "2023-06-01")
 }
 
 func (a *AnthropicMessagesHandlerMethods) parseResponseAndUsage(logger log.Logger, body anthropicMessagesRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {

cmd/cody-gateway/internal/httpapi/completions/fireworks.go

@@ -143,9 +143,13 @@ func (f *FireworksHandlerMethods) getRequestMetadata(body fireworksRequest) (mod
 	return body.Model, map[string]any{"stream": body.Stream}
 }
 
-func (f *FireworksHandlerMethods) transformRequest(r *http.Request) {
-	r.Header.Set("Content-Type", "application/json")
-	r.Header.Set("Authorization", "Bearer "+f.config.AccessToken)
+func (f *FireworksHandlerMethods) transformRequest(downstreamRequest, upstreamRequest *http.Request) {
+	// Enable tracing if the client requests it, see https://readme.fireworks.ai/docs/enabling-tracing
+	if downstreamRequest.Header.Get("X-Fireworks-Genie") == "true" {
+		upstreamRequest.Header.Set("X-Fireworks-Genie", "true")
+	}
+	upstreamRequest.Header.Set("Content-Type", "application/json")
+	upstreamRequest.Header.Set("Authorization", "Bearer "+f.config.AccessToken)
 }
 
 func (f *FireworksHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody fireworksRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {

cmd/cody-gateway/internal/httpapi/completions/google.go

@@ -103,8 +103,8 @@ func (*GoogleHandlerMethods) getRequestMetadata(body googleRequest) (model strin
 	return body.Model, map[string]any{"stream": body.ShouldStream()}
 }
 
-func (o *GoogleHandlerMethods) transformRequest(r *http.Request) {
-	r.Header.Set("Content-Type", "application/json")
+func (o *GoogleHandlerMethods) transformRequest(downstreamRequest, upstreamRequest *http.Request) {
+	upstreamRequest.Header.Set("Content-Type", "application/json")
 }
 
 func (*GoogleHandlerMethods) parseResponseAndUsage(logger log.Logger, reqBody googleRequest, r io.Reader, isStreamRequest bool) (promptUsage, completionUsage usageStats) {

cmd/cody-gateway/internal/httpapi/completions/openai.go

@@ -145,11 +145,11 @@ func (*OpenAIHandlerMethods) getRequestMetadata(body openaiRequest) (model strin
 	return body.Model, map[string]any{"stream": body.Stream}
 }
 
-func (o *OpenAIHandlerMethods) transformRequest(r *http.Request) {
-	r.Header.Set("Content-Type", "application/json")
-	r.Header.Set("Authorization", "Bearer "+o.config.AccessToken)
+func (o *OpenAIHandlerMethods) transformRequest(downstreamRequest, upstreamRequest *http.Request) {
+	upstreamRequest.Header.Set("Content-Type", "application/json")
+	upstreamRequest.Header.Set("Authorization", "Bearer "+o.config.AccessToken)
 	if o.config.OrgID != "" {
-		r.Header.Set("OpenAI-Organization", o.config.OrgID)
+		upstreamRequest.Header.Set("OpenAI-Organization", o.config.OrgID)
 	}
 }
 

cmd/cody-gateway/internal/httpapi/completions/upstream.go

@@ -91,8 +91,9 @@ type upstreamHandlerMethods[ReqT UpstreamRequest] interface {
 	// provided to assist in abuse detection.
 	transformBody(_ *ReqT, identifier string)
 	// transformRequest can be used to modify the HTTP request before it is sent
-	// upstream. To manipulate the body, use transformBody.
-	transformRequest(*http.Request)
+	// upstream. The downstreamRequest parameter is the request sent from the Gateway client.
+	// To manipulate the body, use transformBody.
+	transformRequest(downstreamRequest, upstreamRequest *http.Request)
 	// getRequestMetadata should extract details about the request we are sending
 	// upstream for validation and tracking purposes. Usage data does not need
 	// to be reported here - instead, use parseResponseAndUsage to extract usage,
@@ -158,15 +159,15 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 	// upstreamHandler is the actual HTTP handle that will perform "all of the things"
 	// in order to call the upstream API. e.g. calling the upstreamHandlerMethods in
 	// the correct order, enforcing rate limits and anti-abuse mechanisms, etc.
-	upstreamHandler := func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
+	upstreamHandler := func(w http.ResponseWriter, downstreamRequest *http.Request) {
+		ctx := downstreamRequest.Context()
 		act := actor.FromContext(ctx)
 
 		// TODO: Investigate using actor propagation handler for extracting
 		// this. We had some issues before getting that to work, so for now
 		// just stick with what we've seen working so far.
-		sgActorID := r.Header.Get("X-Sourcegraph-Actor-UID")
-		sgActorAnonymousUID := r.Header.Get("X-Sourcegraph-Actor-Anonymous-UID")
+		sgActorID := downstreamRequest.Header.Get("X-Sourcegraph-Actor-UID")
+		sgActorAnonymousUID := downstreamRequest.Header.Get("X-Sourcegraph-Actor-Anonymous-UID")
 
 		// Build logger for lifecycle of this request with lots of details.
 		logger := act.Logger(sgtrace.Logger(ctx, baseLogger)).With(
@@ -209,7 +210,7 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 
 		// Parse the request body.
 		var body ReqT
-		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+		if err := json.NewDecoder(downstreamRequest.Body).Decode(&body); err != nil {
 			response.JSONError(logger, w, http.StatusBadRequest, errors.Wrap(err, "failed to parse request body"))
 			return
 		}
@@ -295,14 +296,14 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 		}
 
 		// Create a new request to send upstream, making sure we retain the same context.
-		req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(upstreamPayload))
+		upstreamRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(upstreamPayload))
 		if err != nil {
 			response.JSONError(logger, w, http.StatusInternalServerError, errors.Wrap(err, "failed to create request"))
 			return
 		}
 
 		// Run the request transformer.
-		methods.transformRequest(req)
+		methods.transformRequest(downstreamRequest, upstreamRequest)
 
 		// Retrieve metadata from the initial request.
 		model, requestMetadata := methods.getRequestMetadata(body)
@@ -393,11 +394,11 @@ func makeUpstreamHandler[ReqT UpstreamRequest](
 				logger.Error("failed to log event", log.Error(err))
 			}
 		}()
-		resp, err := httpClient.Do(req)
+		resp, err := httpClient.Do(upstreamRequest)
 		if err != nil {
 			// Ignore reporting errors where client disconnected
-			if req.Context().Err() == context.Canceled && errors.Is(err, context.Canceled) {
-				oteltrace.SpanFromContext(req.Context()).
+			if upstreamRequest.Context().Err() == context.Canceled && errors.Is(err, context.Canceled) {
+				oteltrace.SpanFromContext(upstreamRequest.Context()).
 					SetStatus(codes.Error, err.Error())
 				logger.Info("request canceled", log.Error(err))
 				return