feat: add ChunkProcessor interface for streaming chunk processing

Add ChunkProcessor interface that plugins declaring BodyChunked must
implement. The framework validates this at startup and pre-computes
the list of ChunkProcessors per profile. When streaming (no buffering),
each response body chunk is passed through all ChunkProcessors before
being acked to Envoy, enabling in-flight chunk transformation.

Depends on llm-d#169

Signed-off-by: Noy Itzikowitz <nitzikow@redhat.com>

Author: noyitz

docs/proposals/044-plugin-body-mode-capabilities/README.md

@@ -73,12 +73,27 @@ Plugins that don't implement `ResponseBodyRequirement` get a warning:
 INFO  Response plugin does not declare ResponseBodyRequirement, defaulting to BodyFull  profile=default plugin=legacy-plugin/legacy
 ```
 
+### ChunkProcessor Interface
+
+Plugins that declare `BodyChunked` must also implement `ChunkProcessor`:
+
+```go
+type ChunkProcessor interface {
+    ProcessResponseChunk(ctx context.Context, cycleState *CycleState, chunk []byte, isFinal bool) ([]byte, error)
+}
+```
+
+- Validated at startup: declaring `BodyChunked` without implementing `ChunkProcessor` is a fatal error
+- `ProcessResponseChunk` receives each chunk as it arrives and returns (possibly modified) bytes
+- The returned bytes are sent to Envoy as the ack — the ChunkProcessor can transform the data in-flight
+- `ChunkProcessors` are pre-computed per profile and stored on the `Profile` struct
+
 ## Implementation Phases
 
 | Phase | Scope | Status |
 |-------|-------|--------|
-| 1 — Framework | `ResponseBodyMode`, `ResponseBodyRequirement` interface, pre-computation, conditional buffering, tests | **This PR** |
-| 2 — ChunkProcessor | `ChunkProcessor` interface for `BodyChunked` plugins, per-chunk dispatch in server loop | Next PR |
+| 1 — Framework | `ResponseBodyMode`, `ResponseBodyRequirement` interface, pre-computation, conditional buffering, tests | **Done** (PR #169) |
+| 2 — ChunkProcessor | `ChunkProcessor` interface for `BodyChunked` plugins, per-chunk dispatch in server loop | **This PR** |
 | 3 — Plugin declarations | Existing plugins implement `ResponseBodyRequirement` (ODH repo) | Separate PR |
 
 ## Migration

pkg/config/loader/configloader.go

@@ -96,7 +96,10 @@ func LoadConfiguration(configBytes []byte, handle plugin.Handle, processor datas
 		return nil, err
 	}
 
-	computeResponseBuffering(profiles, logger)
+	if err = computeResponseBuffering(profiles, logger); err != nil {
+		logger.Error(err, "failed to compute response buffering requirements")
+		return nil, err
+	}
 
 	return &config.Config{
 		ProfilePicker:  profilePicker,
@@ -309,13 +312,15 @@ func buildPostProcessors(rawConfig *configapi.PluginRefList, handle plugin.Handl
 	return postProcessors, nil
 }
 
-// computeResponseBuffering pre-computes NeedsResponseBuffering for each profile based on the
-// ResponseBodyMode declared by each response plugin. If any response plugin returns BodyFull
-// or doesn't implement ResponseBodyRequirement, the profile needs buffering.
-func computeResponseBuffering(profiles map[string]*requesthandling.Profile, logger logr.Logger) {
+// computeResponseBuffering pre-computes NeedsResponseBuffering and ChunkProcessors for each
+// profile based on the ResponseBodyMode declared by each response plugin. If any response
+// plugin returns BodyFull or doesn't implement ResponseBodyRequirement, the profile needs
+// buffering. Returns an error if a plugin declares BodyChunked but doesn't implement ChunkProcessor.
+func computeResponseBuffering(profiles map[string]*requesthandling.Profile, logger logr.Logger) error {
 	for name, profile := range profiles {
 		needsBuffering := false
 		var bufferingPlugins []string
+		var chunkProcessors []requesthandling.ChunkProcessor
 
 		for _, rp := range profile.ResponsePlugins {
 			mode := requesthandling.BodyFull
@@ -325,18 +330,30 @@ func computeResponseBuffering(profiles map[string]*requesthandling.Profile, logg
 				logger.Info("Response plugin does not declare ResponseBodyRequirement, defaulting to BodyFull",
 					"profile", name, "plugin", rp.TypedName())
 			}
-			if mode == requesthandling.BodyFull {
+
+			switch mode {
+			case requesthandling.BodyFull:
 				needsBuffering = true
 				bufferingPlugins = append(bufferingPlugins, rp.TypedName().Name)
+			case requesthandling.BodyChunked:
+				cp, ok := rp.(requesthandling.ChunkProcessor)
+				if !ok {
+					return fmt.Errorf("plugin %q in profile %q declares BodyChunked but does not implement ChunkProcessor",
+						rp.TypedName().Name, name)
+				}
+				chunkProcessors = append(chunkProcessors, cp)
 			}
 		}
 
 		profile.NeedsResponseBuffering = needsBuffering
+		profile.ChunkProcessors = chunkProcessors
 		logger.Info("Profile response buffering computed",
 			"profile", name,
 			"needsResponseBuffering", needsBuffering,
-			"bufferingPlugins", bufferingPlugins)
+			"bufferingPlugins", bufferingPlugins,
+			"chunkProcessors", len(chunkProcessors))
 	}
+	return nil
 }
 
 // buildModelSelector iterates all built profiles and, for each model-selector plugin found in

pkg/config/loader/response_buffering_test.go

@@ -18,6 +18,7 @@ package loader
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	"sigs.k8s.io/controller-runtime/pkg/log"
@@ -52,6 +53,22 @@ var (
 	_ requesthandling.ResponseBodyRequirement = &fakeResponsePlugin{}
 )
 
+// fakeChunkPlugin declares BodyChunked and implements ChunkProcessor.
+type fakeChunkPlugin struct {
+	fakeResponsePlugin
+}
+
+func (p *fakeChunkPlugin) ProcessResponseChunk(_ context.Context, _ *plugin.CycleState, chunk []byte, _ bool) ([]byte, error) {
+	return chunk, nil
+}
+
+var _ requesthandling.ChunkProcessor = &fakeChunkPlugin{}
+
+// badChunkedPlugin declares BodyChunked but does NOT implement ChunkProcessor.
+type badChunkedPlugin struct {
+	fakeResponsePlugin
+}
+
 type legacyResponsePlugin struct {
 	name string
 }
@@ -71,10 +88,13 @@ func modePtr(m requesthandling.ResponseBodyMode) *requesthandling.ResponseBodyMo
 func TestComputeResponseBuffering(t *testing.T) {
 	logger := log.FromContext(logutil.NewTestLoggerIntoContext(context.Background()))
 
+	chunkedMode := modePtr(requesthandling.BodyChunked)
+
 	tests := []struct {
-		name           string
-		plugins        []requesthandling.ResponseProcessor
-		wantBuffering  bool
+		name              string
+		plugins           []requesthandling.ResponseProcessor
+		wantBuffering     bool
+		wantChunkCount    int
 	}{
 		{
 			name:          "no response plugins",
@@ -90,11 +110,12 @@ func TestComputeResponseBuffering(t *testing.T) {
 			wantBuffering: false,
 		},
 		{
-			name: "all BodyChunked",
+			name: "BodyChunked with ChunkProcessor",
 			plugins: []requesthandling.ResponseProcessor{
-				&fakeResponsePlugin{name: "a", mode: modePtr(requesthandling.BodyChunked)},
+				&fakeChunkPlugin{fakeResponsePlugin{name: "chunker", mode: chunkedMode}},
 			},
-			wantBuffering: false,
+			wantBuffering:  false,
+			wantChunkCount: 1,
 		},
 		{
 			name: "one BodyFull forces buffering",
@@ -114,18 +135,20 @@ func TestComputeResponseBuffering(t *testing.T) {
 		{
 			name: "mixed: BodyChunked + legacy forces buffering",
 			plugins: []requesthandling.ResponseProcessor{
-				&fakeResponsePlugin{name: "a", mode: modePtr(requesthandling.BodyChunked)},
+				&fakeChunkPlugin{fakeResponsePlugin{name: "a", mode: chunkedMode}},
 				&legacyResponsePlugin{name: "legacy"},
 			},
-			wantBuffering: true,
+			wantBuffering:  true,
+			wantChunkCount: 1,
 		},
 		{
 			name: "mixed: BodyNotNeeded + BodyChunked — no buffering",
 			plugins: []requesthandling.ResponseProcessor{
 				&fakeResponsePlugin{name: "a", mode: modePtr(requesthandling.BodyNotNeeded)},
-				&fakeResponsePlugin{name: "b", mode: modePtr(requesthandling.BodyChunked)},
+				&fakeChunkPlugin{fakeResponsePlugin{name: "b", mode: chunkedMode}},
 			},
-			wantBuffering: false,
+			wantBuffering:  false,
+			wantChunkCount: 1,
 		},
 	}
 
@@ -136,35 +159,75 @@ func TestComputeResponseBuffering(t *testing.T) {
 					ResponsePlugins: tc.plugins,
 				},
 			}
-			computeResponseBuffering(profiles, logger)
+			if err := computeResponseBuffering(profiles, logger); err != nil {
+				t.Fatalf("computeResponseBuffering() unexpected error: %v", err)
+			}
 			if profiles["test"].NeedsResponseBuffering != tc.wantBuffering {
 				t.Errorf("NeedsResponseBuffering = %v, want %v", profiles["test"].NeedsResponseBuffering, tc.wantBuffering)
 			}
+			if got := len(profiles["test"].ChunkProcessors); got != tc.wantChunkCount {
+				t.Errorf("ChunkProcessors count = %d, want %d", got, tc.wantChunkCount)
+			}
 		})
 	}
 }
 
+func TestComputeResponseBuffering_BodyChunkedWithoutChunkProcessor(t *testing.T) {
+	logger := log.FromContext(logutil.NewTestLoggerIntoContext(context.Background()))
+
+	profiles := map[string]*requesthandling.Profile{
+		"test": {
+			ResponsePlugins: []requesthandling.ResponseProcessor{
+				&badChunkedPlugin{fakeResponsePlugin{name: "bad", mode: modePtr(requesthandling.BodyChunked)}},
+			},
+		},
+	}
+
+	err := computeResponseBuffering(profiles, logger)
+	if err == nil {
+		t.Fatal("expected error for BodyChunked plugin without ChunkProcessor")
+	}
+	if !strings.Contains(err.Error(), "does not implement ChunkProcessor") {
+		t.Errorf("error message should mention ChunkProcessor, got: %v", err)
+	}
+}
+
 func TestComputeResponseBuffering_MultipleProfiles(t *testing.T) {
 	logger := log.FromContext(logutil.NewTestLoggerIntoContext(context.Background()))
 
+	chunkedMode := modePtr(requesthandling.BodyChunked)
+
 	profiles := map[string]*requesthandling.Profile{
 		"streaming": {
 			ResponsePlugins: []requesthandling.ResponseProcessor{
 				&fakeResponsePlugin{name: "headers-only", mode: modePtr(requesthandling.BodyNotNeeded)},
 			},
 		},
+		"chunked": {
+			ResponsePlugins: []requesthandling.ResponseProcessor{
+				&fakeChunkPlugin{fakeResponsePlugin{name: "meter", mode: chunkedMode}},
+			},
+		},
 		"full-body": {
 			ResponsePlugins: []requesthandling.ResponseProcessor{
 				&fakeResponsePlugin{name: "translator", mode: modePtr(requesthandling.BodyFull)},
 			},
 		},
 	}
 
-	computeResponseBuffering(profiles, logger)
+	if err := computeResponseBuffering(profiles, logger); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
 
 	if profiles["streaming"].NeedsResponseBuffering {
 		t.Error("streaming profile should not need buffering")
 	}
+	if profiles["chunked"].NeedsResponseBuffering {
+		t.Error("chunked profile should not need buffering")
+	}
+	if len(profiles["chunked"].ChunkProcessors) != 1 {
+		t.Errorf("chunked profile should have 1 ChunkProcessor, got %d", len(profiles["chunked"].ChunkProcessors))
+	}
 	if !profiles["full-body"].NeedsResponseBuffering {
 		t.Error("full-body profile should need buffering")
 	}

pkg/framework/interface/requesthandling/plugins.go

@@ -89,3 +89,11 @@ func (m ResponseBodyMode) String() string {
 type ResponseBodyRequirement interface {
 	ResponseBodyMode() ResponseBodyMode
 }
+
+// ChunkProcessor allows a response plugin to process individual response body chunks without
+// waiting for the full body. Plugins declaring BodyChunked MUST implement this interface
+// (validated at startup). The framework calls ProcessResponseChunk for each chunk as it arrives,
+// then acks the chunk to Envoy so it's forwarded to the client immediately.
+type ChunkProcessor interface {
+	ProcessResponseChunk(ctx context.Context, cycleState *plugin.CycleState, chunk []byte, isFinal bool) ([]byte, error)
+}

pkg/framework/interface/requesthandling/types.go

@@ -126,4 +126,9 @@ type Profile struct {
 	// HandleResponseBody with the full body on EndOfStream. When false, each chunk is acked
 	// immediately and forwarded to the client without buffering.
 	NeedsResponseBuffering bool
+
+	// ChunkProcessors are the response plugins that implement ChunkProcessor, pre-computed at
+	// startup. When NeedsResponseBuffering is false, these are called for each response body
+	// chunk before it is acked to Envoy.
+	ChunkProcessors []ChunkProcessor
 }

pkg/handlers/response.go

@@ -131,9 +131,10 @@ func (s *Server) generateEmptyResponseBodyResponse(responseBodyBytes []byte) []*
 	return responses
 }
 
-// ackResponseBodyChunk returns an immediate ack for a response body chunk, allowing Envoy
-// to forward it to the client without waiting for the full body to be accumulated.
-func (s *Server) ackResponseBodyChunk(body *eppb.HttpBody) []*eppb.ProcessingResponse {
+// ackResponseBodyChunkData returns an immediate ack for a response body chunk with the given
+// data and EndOfStream flag, allowing Envoy to forward it to the client without waiting for
+// the full body to be accumulated.
+func (s *Server) ackResponseBodyChunkData(data []byte, endOfStream bool) []*eppb.ProcessingResponse {
 	return []*eppb.ProcessingResponse{
 		{
 			Response: &eppb.ProcessingResponse_ResponseBody{
@@ -142,8 +143,8 @@ func (s *Server) ackResponseBodyChunk(body *eppb.HttpBody) []*eppb.ProcessingRes
 						BodyMutation: &eppb.BodyMutation{
 							Mutation: &eppb.BodyMutation_StreamedResponse{
 								StreamedResponse: &eppb.StreamedBodyResponse{
-									Body:        body.Body,
-									EndOfStream: body.EndOfStream,
+									Body:        data,
+									EndOfStream: endOfStream,
 								},
 							},
 						},
@@ -154,6 +155,21 @@ func (s *Server) ackResponseBodyChunk(body *eppb.HttpBody) []*eppb.ProcessingRes
 	}
 }
 
+// runChunkProcessors runs all ChunkProcessors for the selected profile on a response body chunk.
+func (s *Server) runChunkProcessors(ctx context.Context, reqCtx *RequestContext, chunk []byte, isFinal bool) ([]byte, error) {
+	logger := log.FromContext(ctx).V(logutil.DEFAULT)
+	data := chunk
+	for _, cp := range reqCtx.Profile.ChunkProcessors {
+		var err error
+		data, err = cp.ProcessResponseChunk(ctx, reqCtx.CycleState, data, isFinal)
+		if err != nil {
+			logger.Error(err, "ChunkProcessor failed")
+			return nil, err
+		}
+	}
+	return data, nil
+}
+
 // HandleResponseTrailers handles response trailers.
 func (s *Server) HandleResponseTrailers(trailers *eppb.HttpTrailers) ([]*eppb.ProcessingResponse, error) {
 	return []*eppb.ProcessingResponse{