Repository: envoyproxy/ai-gateway
Pull request #296 (2025-02-06 02:01:57)

extproc: eliminates context.Background
**Commit Message**

This refactors the extproc package and 
eliminates all context.Background usages.


Patch:
diff --git a/internal/controller/controller.go b/internal/controller/controller.go
index b0047a9..6056bdc 100644
--- a/internal/controller/controller.go
+++ b/internal/controller/controller.go
@@ -63,7 +63,7 @@ func StartControllers(ctx context.Context, config *rest.Config, logger logr.Logg
 
 	c := mgr.GetClient()
 	indexer := mgr.GetFieldIndexer()
-	if err = applyIndexing(indexer.IndexField); err != nil {
+	if err = applyIndexing(ctx, indexer.IndexField); err != nil {
 		return fmt.Errorf("failed to apply indexing: %w", err)
 	}
 
@@ -130,18 +130,18 @@ const (
 	k8sClientIndexBackendSecurityPolicyToReferencingAIServiceBackend = "BackendSecurityPolicyToReferencingAIServiceBackend"
 )
 
-func applyIndexing(indexer func(ctx context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error) error {
-	err := indexer(context.Background(), &aigv1a1.AIGatewayRoute{},
+func applyIndexing(ctx context.Context, indexer func(ctx context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error) error {
+	err := indexer(ctx, &aigv1a1.AIGatewayRoute{},
 		k8sClientIndexBackendToReferencingAIGatewayRoute, aiGatewayRouteIndexFunc)
 	if err != nil {
 		return fmt.Errorf("failed to index field for AIGatewayRoute: %w", err)
 	}
-	err = indexer(context.Background(), &aigv1a1.AIServiceBackend{},
+	err = indexer(ctx, &aigv1a1.AIServiceBackend{},
 		k8sClientIndexBackendSecurityPolicyToReferencingAIServiceBackend, aiServiceBackendIndexFunc)
 	if err != nil {
 		return fmt.Errorf("failed to index field for AIServiceBackend: %w", err)
 	}
-	err = indexer(context.Background(), &aigv1a1.BackendSecurityPolicy{},
+	err = indexer(ctx, &aigv1a1.BackendSecurityPolicy{},
 		k8sClientIndexSecretToReferencingBackendSecurityPolicy, backendSecurityPolicyIndexFunc)
 	if err != nil {
 		return fmt.Errorf("failed to index field for BackendSecurityPolicy: %w", err)
diff --git a/internal/controller/sink.go b/internal/controller/sink.go
index 9bf46b9..470b518 100644
--- a/internal/controller/sink.go
+++ b/internal/controller/sink.go
@@ -82,17 +82,17 @@ func newConfigSink(
 	return c
 }
 
-func (c *configSink) backend(namespace, name string) (*aigv1a1.AIServiceBackend, error) {
+func (c *configSink) backend(ctx context.Context, namespace, name string) (*aigv1a1.AIServiceBackend, error) {
 	backend := &aigv1a1.AIServiceBackend{}
-	if err := c.client.Get(context.Background(), client.ObjectKey{Name: name, Namespace: namespace}, backend); err != nil {
+	if err := c.client.Get(ctx, client.ObjectKey{Name: name, Namespace: namespace}, backend); err != nil {
 		return nil, err
 	}
 	return backend, nil
 }
 
-func (c *configSink) backendSecurityPolicy(namespace, name string) (*aigv1a1.BackendSecurityPolicy, error) {
+func (c *configSink) backendSecurityPolicy(ctx context.Context, namespace, name string) (*aigv1a1.BackendSecurityPolicy, error) {
 	backendSecurityPolicy := &aigv1a1.BackendSecurityPolicy{}
-	if err := c.client.Get(context.Background(), client.ObjectKey{Name: name, Namespace: namespace}, backendSecurityPolicy); err != nil {
+	if err := c.client.Get(ctx, client.ObjectKey{Name: name, Namespace: namespace}, backendSecurityPolicy); err != nil {
 		return nil, err
 	}
 	return backendSecurityPolicy, nil
@@ -107,7 +107,7 @@ func (c *configSink) init(ctx context.Context) error {
 				close(c.eventChan)
 				return
 			case event := <-c.eventChan:
-				c.handleEvent(event)
+				c.handleEvent(ctx, event)
 			}
 		}
 	}()
@@ -115,25 +115,25 @@ func (c *configSink) init(ctx context.Context) error {
 }
 
 // handleEvent handles the event received from the controllers in a single goroutine.
-func (c *configSink) handleEvent(event ConfigSinkEvent) {
+func (c *configSink) handleEvent(ctx context.Context, event ConfigSinkEvent) {
 	switch e := event.(type) {
 	case *aigv1a1.AIServiceBackend:
-		c.syncAIServiceBackend(e)
+		c.syncAIServiceBackend(ctx, e)
 	case *aigv1a1.AIGatewayRoute:
-		c.syncAIGatewayRoute(e)
+		c.syncAIGatewayRoute(ctx, e)
 	case *aigv1a1.BackendSecurityPolicy:
-		c.syncBackendSecurityPolicy(e)
+		c.syncBackendSecurityPolicy(ctx, e)
 	case ConfigSinkEventSecretUpdate:
-		c.syncSecret(e.Namespace, e.Name)
+		c.syncSecret(ctx, e.Namespace, e.Name)
 	default:
 		panic(fmt.Sprintf("unexpected event type: %T", e))
 	}
 }
 
-func (c *configSink) syncAIGatewayRoute(aiGatewayRoute *aigv1a1.AIGatewayRoute) {
+func (c *configSink) syncAIGatewayRoute(ctx context.Context, aiGatewayRoute *aigv1a1.AIGatewayRoute) {
 	// Check if the HTTPRouteFilter exists in the namespace.
 	var httpRouteFilter egv1a1.HTTPRouteFilter
-	err := c.client.Get(context.Background(),
+	err := c.client.Get(ctx,
 		client.ObjectKey{Name: hostRewriteHTTPFilterName, Namespace: aiGatewayRoute.Namespace}, &httpRouteFilter)
 	if apierrors.IsNotFound(err) {
 		httpRouteFilter = egv1a1.HTTPRouteFilter{
@@ -149,7 +149,7 @@ func (c *configSink) syncAIGatewayRoute(aiGatewayRoute *aigv1a1.AIGatewayRoute)
 				},
 			},
 		}
-		if err := c.client.Create(context.Background(), &httpRouteFilter); err != nil {
+		if err := c.client.Create(ctx, &httpRouteFilter); err != nil {
 			c.logger.Error(err, "failed to create HTTPRouteFilter", "namespace", aiGatewayRoute.Namespace, "name", hostRewriteHTTPFilterName)
 			return
 		}
@@ -161,7 +161,7 @@ func (c *configSink) syncAIGatewayRoute(aiGatewayRoute *aigv1a1.AIGatewayRoute)
 	// Check if the HTTPRoute exists.
 	c.logger.Info("syncing AIGatewayRoute", "namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name)
 	var httpRoute gwapiv1.HTTPRoute
-	err = c.client.Get(context.Background(), client.ObjectKey{Name: aiGatewayRoute.Name, Namespace: aiGatewayRoute.Namespace}, &httpRoute)
+	err = c.client.Get(ctx, client.ObjectKey{Name: aiGatewayRoute.Name, Namespace: aiGatewayRoute.Namespace}, &httpRoute)
 	existingRoute := err == nil
 	if apierrors.IsNotFound(err) {
 		// This means that this AIGatewayRoute is a new one.
@@ -181,20 +181,20 @@ func (c *configSink) syncAIGatewayRoute(aiGatewayRoute *aigv1a1.AIGatewayRoute)
 	}
 
 	// Update the HTTPRoute with the new AIGatewayRoute.
-	if err := c.newHTTPRoute(&httpRoute, aiGatewayRoute); err != nil {
+	if err := c.newHTTPRoute(ctx, &httpRoute, aiGatewayRoute); err != nil {
 		c.logger.Error(err, "failed to update HTTPRoute with AIGatewayRoute", "namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name)
 		return
 	}
 
 	if existingRoute {
 		c.logger.Info("updating HTTPRoute", "namespace", httpRoute.Namespace, "name", httpRoute.Name)
-		if err := c.client.Update(context.Background(), &httpRoute); err != nil {
+		if err := c.client.Update(ctx, &httpRoute); err != nil {
 			c.logger.Error(err, "failed to update HTTPRoute", "namespace", httpRoute.Namespace, "name", httpRoute.Name)
 			return
 		}
 	} else {
 		c.logger.Info("creating HTTPRoute", "namespace", httpRoute.Namespace, "name", httpRoute.Name)
-		if err := c.client.Create(context.Background(), &httpRoute); err != nil {
+		if err := c.client.Create(ctx, &httpRoute); err != nil {
 			c.logger.Error(err, "failed to create HTTPRoute", "namespace", httpRoute.Namespace, "name", httpRoute.Name)
 			return
 		}
@@ -202,30 +202,30 @@ func (c *configSink) syncAIGatewayRoute(aiGatewayRoute *aigv1a1.AIGatewayRoute)
 
 	// Update the extproc configmap.
 	uuid := string(uuid2.NewUUID())
-	if err := c.updateExtProcConfigMap(aiGatewayRoute, uuid); err != nil {
+	if err := c.updateExtProcConfigMap(ctx, aiGatewayRoute, uuid); err != nil {
 		c.logger.Error(err, "failed to update extproc configmap", "namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name)
 		return
 	}
 
 	// Deploy extproc deployment with potential updates.
-	err = c.syncExtProcDeployment(context.Background(), aiGatewayRoute)
+	err = c.syncExtProcDeployment(ctx, aiGatewayRoute)
 	if err != nil {
 		c.logger.Error(err, "failed to deploy ext proc", "namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name)
 		return
 	}
 
 	// Annotate all pods with the new config.
-	err = c.annotateExtProcPods(context.Background(), aiGatewayRoute, uuid)
+	err = c.annotateExtProcPods(ctx, aiGatewayRoute, uuid)
 	if err != nil {
 		c.logger.Error(err, "failed to annotate pods", "namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name)
 		return
 	}
 }
 
-func (c *configSink) syncAIServiceBackend(aiBackend *aigv1a1.AIServiceBackend) {
+func (c *configSink) syncAIServiceBackend(ctx context.Context, aiBackend *aigv1a1.AIServiceBackend) {
 	key := fmt.Sprintf("%s.%s", aiBackend.Name, aiBackend.Namespace)
 	var aiGatewayRoutes aigv1a1.AIGatewayRouteList
-	err := c.client.List(context.Background(), &aiGatewayRoutes, client.MatchingFields{k8sClientIndexBackendToReferencingAIGatewayRoute: key})
+	err := c.client.List(ctx, &aiGatewayRoutes, client.MatchingFields{k8sClientIndexBackendToReferencingAIGatewayRoute: key})
 	if err != nil {
 		c.logger.Error(err, "failed to list AIGatewayRoute", "backend", key)
 		return
@@ -235,27 +235,27 @@ func (c *configSink) syncAIServiceBackend(aiBackend *aigv1a1.AIServiceBackend) {
 			"namespace", aiGatewayRoute.Namespace, "name", aiGatewayRoute.Name,
 			"referenced_backend", aiBackend.Name, "referenced_backend_namespace", aiBackend.Namespace,
 		)
-		c.syncAIGatewayRoute(&aiGatewayRoute)
+		c.syncAIGatewayRoute(ctx, &aiGatewayRoute)
 	}
 }
 
-func (c *configSink) syncBackendSecurityPolicy(bsp *aigv1a1.BackendSecurityPolicy) {
+func (c *configSink) syncBackendSecurityPolicy(ctx context.Context, bsp *aigv1a1.BackendSecurityPolicy) {
 	key := fmt.Sprintf("%s.%s", bsp.Name, bsp.Namespace)
 	var aiServiceBackends aigv1a1.AIServiceBackendList
-	err := c.client.List(context.Background(), &aiServiceBackends, client.MatchingFields{k8sClientIndexBackendSecurityPolicyToReferencingAIServiceBackend: key})
+	err := c.client.List(ctx, &aiServiceBackends, client.MatchingFields{k8sClientIndexBackendSecurityPolicyToReferencingAIServiceBackend: key})
 	if err != nil {
 		c.logger.Error(err, "failed to list AIServiceBackendList", "backendSecurityPolicy", key)
 		return
 	}
 	for i := range aiServiceBackends.Items {
 		aiBackend := &aiServiceBackends.Items[i]
-		c.syncAIServiceBackend(aiBackend)
+		c.syncAIServiceBackend(ctx, aiBackend)
 	}
 }
 
 // updateExtProcConfigMap updates the external process configmap with the new AIGatewayRoute.
-func (c *configSink) updateExtProcConfigMap(aiGatewayRoute *aigv1a1.AIGatewayRoute, uuid string) error {
-	configMap, err := c.kube.CoreV1().ConfigMaps(aiGatewayRoute.Namespace).Get(context.Background(), extProcName(aiGatewayRoute), metav1.GetOptions{})
+func (c *configSink) updateExtProcConfigMap(ctx context.Context, aiGatewayRoute *aigv1a1.AIGatewayRoute, uuid string) error {
+	configMap, err := c.kube.CoreV1().ConfigMaps(aiGatewayRoute.Namespace).Get(ctx, extProcName(aiGatewayRoute), metav1.GetOptions{})
 	if err != nil {
 		// This is a bug since we should have created the configmap before sending the AIGatewayRoute to the configSink.
 		panic(fmt.Errorf("failed to get configmap %s: %w", extProcName(aiGatewayRoute), err))
@@ -277,7 +277,7 @@ func (c *configSink) updateExtProcConfigMap(aiGatewayRoute *aigv1a1.AIGatewayRou
 			key := fmt.Sprintf("%s.%s", backend.Name, aiGatewayRoute.Namespace)
 			ec.Rules[i].Backends[j].Name = key
 			ec.Rules[i].Backends[j].Weight = backend.Weight
-			backendObj, err := c.backend(aiGatewayRoute.Namespace, backend.Name)
+			backendObj, err := c.backend(ctx, aiGatewayRoute.Namespace, backend.Name)
 			if err != nil {
 				return fmt.Errorf("failed to get AIServiceBackend %s: %w", key, err)
 			} else {
@@ -289,7 +289,7 @@ func (c *configSink) updateExtProcConfigMap(aiGatewayRoute *aigv1a1.AIGatewayRou
 				volumeName := backendSecurityPolicyVolumeName(
 					i, j, string(backendObj.Spec.BackendSecurityPolicyRef.Name),
 				)
-				backendSecurityPolicy, err := c.backendSecurityPolicy(aiGatewayRoute.Namespace, string(bspRef.Name))
+				backendSecurityPolicy, err := c.backendSecurityPolicy(ctx, aiGatewayRoute.Namespace, string(bspRef.Name))
 				if err != nil {
 					return fmt.Errorf("failed to get BackendSecurityPolicy %s: %w", bspRef.Name, err)
 				}
@@ -357,14 +357,14 @@ func (c *configSink) updateExtProcConfigMap(aiGatewayRoute *aigv1a1.AIGatewayRou
 		configMap.Data = make(map[string]string)
 	}
 	configMap.Data[expProcConfigFileName] = string(marshaled)
-	if _, err := c.kube.CoreV1().ConfigMaps(aiGatewayRoute.Namespace).Update(context.Background(), configMap, metav1.UpdateOptions{}); err != nil {
+	if _, err := c.kube.CoreV1().ConfigMaps(aiGatewayRoute.Namespace).Update(ctx, configMap, metav1.UpdateOptions{}); err != nil {
 		return fmt.Errorf("failed to update configmap %s: %w", configMap.Name, err)
 	}
 	return nil
 }
 
 // newHTTPRoute updates the HTTPRoute with the new AIGatewayRoute.
-func (c *configSink) newHTTPRoute(dst *gwapiv1.HTTPRoute, aiGatewayRoute *aigv1a1.AIGatewayRoute) error {
+func (c *configSink) newHTTPRoute(ctx context.Context, dst *gwapiv1.HTTPRoute, aiGatewayRoute *aigv1a1.AIGatewayRoute) error {
 	var backends []*aigv1a1.AIServiceBackend
 	dedup := make(map[string]struct{})
 	for _, rule := range aiGatewayRoute.Spec.Rules {
@@ -374,7 +374,7 @@ func (c *configSink) newHTTPRoute(dst *gwapiv1.HTTPRoute, aiGatewayRoute *aigv1a
 				continue
 			}
 			dedup[key] = struct{}{}
-			backend, err := c.backend(aiGatewayRoute.Namespace, br.Name)
+			backend, err := c.backend(ctx, aiGatewayRoute.Namespace, br.Name)
 			if err != nil {
 				return fmt.Errorf("AIServiceBackend %s not found", key)
 			}
@@ -516,7 +516,7 @@ func (c *configSink) syncExtProcDeployment(ctx context.Context, aiGatewayRoute *
 			if err := ctrlutil.SetControllerReference(aiGatewayRoute, deployment, c.client.Scheme()); err != nil {
 				panic(fmt.Errorf("BUG: failed to set controller reference for deployment: %w", err))
 			}
-			updatedSpec, err := c.mountBackendSecurityPolicySecrets(&deployment.Spec.Template.Spec, aiGatewayRoute)
+			updatedSpec, err := c.mountBackendSecurityPolicySecrets(ctx, &deployment.Spec.Template.Spec, aiGatewayRoute)
 			if err == nil {
 				deployment.Spec.Template.Spec = *updatedSpec
 			}
@@ -530,7 +530,7 @@ func (c *configSink) syncExtProcDeployment(ctx context.Context, aiGatewayRoute *
 			return fmt.Errorf("failed to get deployment: %w", err)
 		}
 	} else {
-		updatedSpec, err := c.mountBackendSecurityPolicySecrets(&deployment.Spec.Template.Spec, aiGatewayRoute)
+		updatedSpec, err := c.mountBackendSecurityPolicySecrets(ctx, &deployment.Spec.Template.Spec, aiGatewayRoute)
 		if err == nil {
 			deployment.Spec.Template.Spec = *updatedSpec
 		}
@@ -569,7 +569,7 @@ func (c *configSink) syncExtProcDeployment(ctx context.Context, aiGatewayRoute *
 }
 
 // mountBackendSecurityPolicySecrets will mount secrets based on backendSecurityPolicies attached to AIServiceBackend.
-func (c *configSink) mountBackendSecurityPolicySecrets(spec *corev1.PodSpec, aiGatewayRoute *aigv1a1.AIGatewayRoute) (*corev1.PodSpec, error) {
+func (c *configSink) mountBackendSecurityPolicySecrets(ctx context.Context, spec *corev1.PodSpec, aiGatewayRoute *aigv1a1.AIGatewayRoute) (*corev1.PodSpec, error) {
 	// Mount from scratch to avoid secrets that should be unmounted.
 	// Only keep the original mount which should be the config volume.
 	spec.Volumes = spec.Volumes[:1]
@@ -580,13 +580,13 @@ func (c *configSink) mountBackendSecurityPolicySecrets(spec *corev1.PodSpec, aiG
 		rule := &aiGatewayRoute.Spec.Rules[i]
 		for j := range rule.BackendRefs {
 			backendRef := &rule.BackendRefs[j]
-			backend, err := c.backend(aiGatewayRoute.Namespace, backendRef.Name)
+			backend, err := c.backend(ctx, aiGatewayRoute.Namespace, backendRef.Name)
 			if err != nil {
 				return nil, fmt.Errorf("failed to get backend %s: %w", backendRef.Name, err)
 			}
 
 			if backendSecurityPolicyRef := backend.Spec.BackendSecurityPolicyRef; backendSecurityPolicyRef != nil {
-				backendSecurityPolicy, err := c.backendSecurityPolicy(aiGatewayRoute.Namespace, string(backendSecurityPolicyRef.Name))
+				backendSecurityPolicy, err := c.backendSecurityPolicy(ctx, aiGatewayRoute.Namespace, string(backendSecurityPolicyRef.Name))
 				if err != nil {
 					return nil, fmt.Errorf("failed to get backend security policy %s: %w", backendSecurityPolicyRef.Name, err)
 				}
@@ -626,9 +626,9 @@ func (c *configSink) mountBackendSecurityPolicySecrets(spec *corev1.PodSpec, aiG
 }
 
 // syncSecret syncs the state of all resource referencing the given secret.
-func (c *configSink) syncSecret(namespace, name string) {
+func (c *configSink) syncSecret(ctx context.Context, namespace, name string) {
 	var backendSecurityPolicies aigv1a1.BackendSecurityPolicyList
-	err := c.client.List(context.Background(), &backendSecurityPolicies,
+	err := c.client.List(ctx, &backendSecurityPolicies,
 		client.MatchingFields{
 			k8sClientIndexSecretToReferencingBackendSecurityPolicy: fmt.Sprintf("%s.%s", name, namespace),
 		},
@@ -639,7 +639,7 @@ func (c *configSink) syncSecret(namespace, name string) {
 	}
 	for i := range backendSecurityPolicies.Items {
 		backendSecurityPolicy := &backendSecurityPolicies.Items[i]
-		c.syncBackendSecurityPolicy(backendSecurityPolicy)
+		c.syncBackendSecurityPolicy(ctx, backendSecurityPolicy)
 	}
 }
 
diff --git a/internal/extproc/backendauth/api_key.go b/internal/extproc/backendauth/api_key.go
index 3a94fb0..e7c3686 100644
--- a/internal/extproc/backendauth/api_key.go
+++ b/internal/extproc/backendauth/api_key.go
@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -27,7 +28,7 @@ func newAPIKeyHandler(auth *filterapi.APIKeyAuth) (Handler, error) {
 // Do implements [Handler.Do].
 //
 // Extracts the api key from the local file and set it as an authorization header.
-func (a *apiKeyHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
+func (a *apiKeyHandler) Do(_ context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, _ *extprocv3.BodyMutation) error {
 	requestHeaders["Authorization"] = fmt.Sprintf("Bearer %s", a.apiKey)
 	headerMut.SetHeaders = append(headerMut.SetHeaders, &corev3.HeaderValueOption{
 		Header: &corev3.HeaderValue{Key: "Authorization", RawValue: []byte(requestHeaders["Authorization"])},
diff --git a/internal/extproc/backendauth/auth.go b/internal/extproc/backendauth/auth.go
index e837584..8dd4fbc 100644
--- a/internal/extproc/backendauth/auth.go
+++ b/internal/extproc/backendauth/auth.go
@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"errors"
 
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -13,7 +14,7 @@ import (
 // TODO: maybe this can be just "post-transformation" handler, as it is not really only about auth.
 type Handler interface {
 	// Do performs the backend auth, and make changes to the request headers and body mutations.
-	Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error
+	Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error
 }
 
 // NewHandler returns a new implementation of [Handler] based on the configuration.
diff --git a/internal/extproc/backendauth/aws.go b/internal/extproc/backendauth/aws.go
index 65fbf09..e01ecc6 100644
--- a/internal/extproc/backendauth/aws.go
+++ b/internal/extproc/backendauth/aws.go
@@ -27,7 +27,7 @@ type awsHandler struct {
 	region      string
 }
 
-func newAWSHandler(awsAuth *filterapi.AWSAuth) (*awsHandler, error) {
+func newAWSHandler(awsAuth *filterapi.AWSAuth) (Handler, error) {
 	var credentials aws.Credentials
 	var region string
 
@@ -60,7 +60,7 @@ func newAWSHandler(awsAuth *filterapi.AWSAuth) (*awsHandler, error) {
 //
 // This assumes that during the transformation, the path is set in the header mutation as well as
 // the body in the body mutation.
-func (a *awsHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error {
+func (a *awsHandler) Do(ctx context.Context, requestHeaders map[string]string, headerMut *extprocv3.HeaderMutation, bodyMut *extprocv3.BodyMutation) error {
 	method := requestHeaders[":method"]
 	path := ""
 	if headerMut.SetHeaders != nil {
@@ -90,7 +90,7 @@ func (a *awsHandler) Do(requestHeaders map[string]string, headerMut *extprocv3.H
 		return fmt.Errorf("cannot create request: %w", err)
 	}
 
-	err = a.signer.SignHTTP(context.Background(), a.credentials, req,
+	err = a.signer.SignHTTP(ctx, a.credentials, req,
 		hex.EncodeToString(payloadHash[:]), "bedrock", a.region, time.Now())
 	if err != nil {
 		return fmt.Errorf("cannot sign request: %w", err)
diff --git a/internal/extproc/processor.go b/internal/extproc/processor.go
index 7af7842..bf1610e 100644
--- a/internal/extproc/processor.go
+++ b/internal/extproc/processor.go
@@ -80,7 +80,7 @@ func (p *Processor) ProcessRequestHeaders(_ context.Context, headers *corev3.Hea
 }
 
 // ProcessRequestBody implements [Processor.ProcessRequestBody].
-func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
+func (p *Processor) ProcessRequestBody(ctx context.Context, rawBody *extprocv3.HttpBody) (res *extprocv3.ProcessingResponse, err error) {
 	path := p.requestHeaders[":path"]
 	model, body, err := p.config.bodyParser(path, rawBody)
 	if err != nil {
@@ -122,7 +122,7 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt
 	})
 
 	if authHandler, ok := p.config.backendAuthHandlers[b.Name]; ok {
-		if err := authHandler.Do(p.requestHeaders, headerMutation, bodyMutation); err != nil {
+		if err := authHandler.Do(ctx, p.requestHeaders, headerMutation, bodyMutation); err != nil {
 			return nil, fmt.Errorf("failed to do auth request: %w", err)
 		}
 	}


Test patch:
diff --git a/internal/controller/sink_test.go b/internal/controller/sink_test.go
index 3cdd3cb..df0aae8 100644
--- a/internal/controller/sink_test.go
+++ b/internal/controller/sink_test.go
@@ -30,7 +30,7 @@ import (
 
 func requireNewFakeClientWithIndexes(t *testing.T) client.Client {
 	builder := fake.NewClientBuilder().WithScheme(scheme)
-	err := applyIndexing(func(ctx context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error {
+	err := applyIndexing(context.Background(), func(ctx context.Context, obj client.Object, field string, extractValue client.IndexerFunc) error {
 		builder = builder.WithIndex(obj, field, extractValue)
 		return nil
 	})
@@ -51,10 +51,10 @@ func TestConfigSink_handleEvent(t *testing.T) {
 	}, metav1.CreateOptions{})
 	require.NoError(t, err)
 
-	s.handleEvent(ConfigSinkEventSecretUpdate{Namespace: "ns", Name: "some-secret"})
-	s.handleEvent(&aigv1a1.AIServiceBackend{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
-	s.handleEvent(&aigv1a1.BackendSecurityPolicy{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
-	s.handleEvent(&aigv1a1.AIGatewayRoute{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
+	s.handleEvent(context.Background(), ConfigSinkEventSecretUpdate{Namespace: "ns", Name: "some-secret"})
+	s.handleEvent(context.Background(), &aigv1a1.AIServiceBackend{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
+	s.handleEvent(context.Background(), &aigv1a1.BackendSecurityPolicy{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
+	s.handleEvent(context.Background(), &aigv1a1.AIGatewayRoute{ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"}})
 }
 
 func TestConfigSink_syncAIGatewayRoute(t *testing.T) {
@@ -105,7 +105,7 @@ func TestConfigSink_syncAIGatewayRoute(t *testing.T) {
 		require.NoError(t, err)
 
 		// Then sync, which should update the HTTPRoute.
-		s.syncAIGatewayRoute(route)
+		s.syncAIGatewayRoute(context.Background(), route)
 		var updatedHTTPRoute gwapiv1.HTTPRoute
 		err = fakeClient.Get(context.Background(), client.ObjectKey{Name: "route1", Namespace: "ns1"}, &updatedHTTPRoute)
 		require.NoError(t, err)
@@ -144,7 +144,7 @@ func TestConfigSink_syncAIServiceBackend(t *testing.T) {
 	require.NoError(t, fakeClient.Create(context.Background(), route, &client.CreateOptions{}))
 
 	s := newConfigSink(fakeClient, nil, logr.Discard(), eventChan, "defaultExtProcImage", "debug")
-	s.syncAIServiceBackend(&aigv1a1.AIServiceBackend{
+	s.syncAIServiceBackend(context.Background(), &aigv1a1.AIServiceBackend{
 		ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns1"},
 		Spec: aigv1a1.AIServiceBackendSpec{
 			BackendRef: gwapiv1.BackendObjectReference{Name: "some-backend", Namespace: ptr.To[gwapiv1.Namespace]("ns1")},
@@ -166,7 +166,7 @@ func TestConfigSink_syncBackendSecurityPolicy(t *testing.T) {
 	require.NoError(t, fakeClient.Create(context.Background(), &backend, &client.CreateOptions{}))
 
 	s := newConfigSink(fakeClient, nil, logr.Discard(), eventChan, "defaultExtProcImage", "debug")
-	s.syncBackendSecurityPolicy(&aigv1a1.BackendSecurityPolicy{
+	s.syncBackendSecurityPolicy(context.Background(), &aigv1a1.BackendSecurityPolicy{
 		ObjectMeta: metav1.ObjectMeta{Name: "apple", Namespace: "ns"},
 	})
 }
@@ -235,7 +235,7 @@ func Test_newHTTPRoute(t *testing.T) {
 		err := s.client.Create(context.Background(), backend, &client.CreateOptions{})
 		require.NoError(t, err)
 	}
-	err := s.newHTTPRoute(httpRoute, aiGatewayRoute)
+	err := s.newHTTPRoute(context.Background(), httpRoute, aiGatewayRoute)
 	require.NoError(t, err)
 
 	expRules := []gwapiv1.HTTPRouteRule{
@@ -466,7 +466,7 @@ func Test_updateExtProcConfigMap(t *testing.T) {
 			}, metav1.CreateOptions{})
 			require.NoError(t, err)
 
-			err = s.updateExtProcConfigMap(tc.route, tc.exp.UUID)
+			err = s.updateExtProcConfigMap(context.Background(), tc.route, tc.exp.UUID)
 			require.NoError(t, err)
 
 			cm, err := s.kube.CoreV1().ConfigMaps(tc.route.Namespace).Get(context.Background(), extProcName(tc.route), metav1.GetOptions{})
@@ -776,7 +776,7 @@ func TestConfigSink_MountBackendSecurityPolicySecrets(t *testing.T) {
 
 	require.NoError(t, fakeClient.Create(context.Background(), &aiGateway, &client.CreateOptions{}))
 
-	updatedSpec, err := s.mountBackendSecurityPolicySecrets(&spec, &aiGateway)
+	updatedSpec, err := s.mountBackendSecurityPolicySecrets(context.Background(), &spec, &aiGateway)
 	require.NoError(t, err)
 
 	require.Len(t, updatedSpec.Volumes, 3)
@@ -809,7 +809,7 @@ func TestConfigSink_MountBackendSecurityPolicySecrets(t *testing.T) {
 	require.NoError(t, fakeClient.Create(context.Background(), &backend, &client.CreateOptions{}))
 	require.NotNil(t, s)
 
-	updatedSpec, err = s.mountBackendSecurityPolicySecrets(&spec, &aiGateway)
+	updatedSpec, err = s.mountBackendSecurityPolicySecrets(context.Background(), &spec, &aiGateway)
 	require.NoError(t, err)
 
 	require.Len(t, updatedSpec.Volumes, 3)
@@ -877,5 +877,5 @@ func Test_syncSecret(t *testing.T) {
 		Data:       map[string][]byte{"key": []byte("value")},
 	}, metav1.CreateOptions{})
 	require.NoError(t, err)
-	s.syncSecret("ns", "some-secret")
+	s.syncSecret(context.Background(), "ns", "some-secret")
 }
diff --git a/internal/extproc/backendauth/api_key_test.go b/internal/extproc/backendauth/api_key_test.go
index cfb3ce1..6f08f38 100644
--- a/internal/extproc/backendauth/api_key_test.go
+++ b/internal/extproc/backendauth/api_key_test.go
@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -62,7 +63,7 @@ func TestApiKeyHandler_Do(t *testing.T) {
 			Body: []byte(`{"messages": [{"role": "user", "content": [{"text": "Say this is a test!"}]}]}`),
 		},
 	}
-	err = handler.Do(requestHeaders, headerMut, bodyMut)
+	err = handler.Do(context.Background(), requestHeaders, headerMut, bodyMut)
 	require.NoError(t, err)
 
 	bearerToken, ok := requestHeaders["Authorization"]
diff --git a/internal/extproc/backendauth/aws_test.go b/internal/extproc/backendauth/aws_test.go
index 3a3da56..e0d922b 100644
--- a/internal/extproc/backendauth/aws_test.go
+++ b/internal/extproc/backendauth/aws_test.go
@@ -1,6 +1,7 @@
 package backendauth
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -42,7 +43,7 @@ func TestAWSHandler_Do(t *testing.T) {
 
 	for _, tc := range []struct {
 		name    string
-		handler *awsHandler
+		handler Handler
 	}{
 		{
 			name:    "Using AWS Credential File",
@@ -64,7 +65,7 @@ func TestAWSHandler_Do(t *testing.T) {
 					Body: []byte(`{"messages": [{"role": "user", "content": [{"text": "Say this is a test!"}]}]}`),
 				},
 			}
-			err = tc.handler.Do(requestHeaders, headerMut, bodyMut)
+			err = tc.handler.Do(context.Background(), requestHeaders, headerMut, bodyMut)
 			require.NoError(t, err)
 		})
 	}