Repository: envoyproxy/ai-gateway
Pull request #482 (2025-03-12 16:47:15)

extproc: make Main function accept args, stderr
**Commit Message**

This makes changes to Main function of the extproc so that it will accept context as well as args, etc to allow the callsite to fine tune the behavior of extproc.

**Related Issues/PRs (if applicable)**

Extracted from #458 




Patch:
diff --git a/cmd/extproc/main.go b/cmd/extproc/main.go
index f91b7c8..937d10a 100644
--- a/cmd/extproc/main.go
+++ b/cmd/extproc/main.go
@@ -5,6 +5,22 @@
 
 package main
 
-import "github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
 
-func main() { mainlib.Main() }
+	"github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
+)
+
+func main() {
+	ctx, cancel := context.WithCancel(context.Background())
+	signalsChan := make(chan os.Signal, 1)
+	signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-signalsChan
+		cancel()
+	}()
+	mainlib.Main(ctx, os.Args[1:], os.Stderr)
+}
diff --git a/cmd/extproc/mainlib/main.go b/cmd/extproc/mainlib/main.go
index b5f4fbd..7a29c79 100644
--- a/cmd/extproc/mainlib/main.go
+++ b/cmd/extproc/mainlib/main.go
@@ -10,14 +10,12 @@ import (
 	"errors"
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"log/slog"
 	"net"
 	"net/http"
-	"os"
-	"os/signal"
 	"strings"
-	"syscall"
 	"time"
 
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -84,13 +82,17 @@ func parseAndValidateFlags(args []string) (extProcFlags, error) {
 
 // Main is a main function for the external processor exposed
 // for allowing users to build their own external processor.
-func Main() {
-	flags, err := parseAndValidateFlags(os.Args[1:])
+//
+// * ctx is the context for the external processor.
+// * args are the command line arguments passed to the external processor without the program name.
+// * stderr is the writer to use for standard error where the external processor will output logs.
+func Main(ctx context.Context, args []string, stderr io.Writer) {
+	flags, err := parseAndValidateFlags(args)
 	if err != nil {
 		log.Fatalf("failed to parse and validate extProcFlags: %v", err)
 	}
 
-	l := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: flags.logLevel}))
+	l := slog.New(slog.NewTextHandler(stderr, &slog.HandlerOptions{Level: flags.logLevel}))
 
 	l.Info("starting external processor",
 		slog.String("version", version.Version),
@@ -98,14 +100,6 @@ func Main() {
 		slog.String("configPath", flags.configPath),
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
-	signalsChan := make(chan os.Signal, 1)
-	signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
-	go func() {
-		<-signalsChan
-		cancel()
-	}()
-
 	lis, err := net.Listen(listenAddress(flags.extProcAddr))
 	if err != nil {
 		log.Fatalf("failed to listen: %v", err)
diff --git a/examples/extproc_custom_router/main.go b/examples/extproc_custom_router/main.go
index 11e4281..d6ba594 100644
--- a/examples/extproc_custom_router/main.go
+++ b/examples/extproc_custom_router/main.go
@@ -6,7 +6,11 @@
 package main
 
 import (
+	"context"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/envoyproxy/ai-gateway/cmd/extproc/mainlib"
 	"github.com/envoyproxy/ai-gateway/filterapi"
@@ -42,5 +46,12 @@ func main() {
 	// Initializes the custom router.
 	x.NewCustomRouter = newCustomRouter
 	// Executes the main function of the external processor.
-	mainlib.Main()
+	ctx, cancel := context.WithCancel(context.Background())
+	signalsChan := make(chan os.Signal, 1)
+	signal.Notify(signalsChan, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-signalsChan
+		cancel()
+	}()
+	mainlib.Main(ctx, os.Args[1:], os.Stderr)
 }
diff --git a/internal/extproc/chatcompletion_processor.go b/internal/extproc/chatcompletion_processor.go
index 0bdb2be..3407ec5 100644
--- a/internal/extproc/chatcompletion_processor.go
+++ b/internal/extproc/chatcompletion_processor.go
@@ -99,7 +99,7 @@ func (c *chatCompletionProcessor) ProcessRequestBody(ctx context.Context, rawBod
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse request body: %w", err)
 	}
-	c.logger.Info("Processing request", "path", c.requestHeaders[":path"], "model", model)
+	c.logger.Info("processing request body", "path", c.requestHeaders[":path"], "model", model)
 
 	c.metrics.SetModel(model)
 	c.requestHeaders[c.config.modelNameHeaderKey] = model
diff --git a/internal/extproc/server.go b/internal/extproc/server.go
index 5d287c9..56bcc5d 100644
--- a/internal/extproc/server.go
+++ b/internal/extproc/server.go
@@ -276,7 +276,11 @@ func filterSensitiveBodyForLogging(resp *extprocv3.ProcessingResponse, logger *s
 	if resp == nil {
 		return &extprocv3.ProcessingResponse{}
 	}
-	original := resp.Response.(*extprocv3.ProcessingResponse_RequestBody)
+	original, ok := resp.Response.(*extprocv3.ProcessingResponse_RequestBody)
+	if !ok {
+		// Meaning this is the immediate response, that doesn't need to be filtered.
+		return resp
+	}
 	originalHeaderMutation := original.RequestBody.Response.GetHeaderMutation()
 	redactedHeaderMutation := &extprocv3.HeaderMutation{
 		RemoveHeaders: originalHeaderMutation.GetRemoveHeaders(),


Test patch:
diff --git a/internal/extproc/server_test.go b/internal/extproc/server_test.go
index 241a644..ea5f407 100644
--- a/internal/extproc/server_test.go
+++ b/internal/extproc/server_test.go
@@ -380,6 +380,17 @@ func Test_filterSensitiveBodyForLogging(t *testing.T) {
 		{Header: &corev3.HeaderValue{Key: "Authorization", RawValue: []byte("sensitive")}},
 	}, originalMutation.GetSetHeaders())
 	require.Contains(t, buf.String(), "filtering sensitive header")
+
+	t.Run("do nothing for immediate response", func(t *testing.T) {
+		resp := &extprocv3.ProcessingResponse{
+			Response: &extprocv3.ProcessingResponse_ImmediateResponse{
+				ImmediateResponse: &extprocv3.ImmediateResponse{},
+			},
+		}
+		filtered := filterSensitiveBodyForLogging(resp, logger, []string{"authorization"})
+		require.NotNil(t, filtered)
+		require.Equal(t, resp, filtered)
+	})
 }
 
 func Test_headersToMap(t *testing.T) {