Merge remote-tracking branch 'upstream/main' into rhoai-3.5

Author: m-rafeeq

deploy/payload-processing/values.yaml

@@ -21,6 +21,7 @@ upstreamIpp:
           fieldName: model
           headerName: X-Gateway-Model-Name
       - type: model-provider-resolver
+      - type: stream-usage-enforcer
       - type: api-translation
       - type: apikey-injection
       profiles:
@@ -30,6 +31,7 @@ upstreamIpp:
           - pluginRef: maas-headers-guard
           - pluginRef: model-extractor
           - pluginRef: model-provider-resolver
+          - pluginRef: stream-usage-enforcer
           - pluginRef: api-translation
           - pluginRef: apikey-injection
           response:

pkg/plugins/plugins.go

@@ -22,6 +22,7 @@ import (
 	maas_headers_guard "github.com/opendatahub-io/ai-gateway-payload-processing/pkg/plugins/maas-headers-guard"
 	provider_resolver "github.com/opendatahub-io/ai-gateway-payload-processing/pkg/plugins/model-provider-resolver"
 	"github.com/opendatahub-io/ai-gateway-payload-processing/pkg/plugins/nemo"
+	stream_usage_enforcer "github.com/opendatahub-io/ai-gateway-payload-processing/pkg/plugins/stream-usage-enforcer"
 	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin"
 )
 
@@ -32,4 +33,5 @@ func RegisterPlugins() {
 	plugin.Register(apikey_injection.APIKeyInjectionPluginType, apikey_injection.APIKeyInjectionFactory)
 	plugin.Register(nemo.NemoRequestGuardPluginType, nemo.NemoRequestGuardFactory)
 	plugin.Register(nemo.NemoResponseGuardPluginType, nemo.NemoResponseGuardFactory)
+	plugin.Register(stream_usage_enforcer.PluginType, stream_usage_enforcer.Factory)
 }

pkg/plugins/stream-usage-enforcer/plugin.go

@@ -0,0 +1,74 @@
+/*
+Copyright 2026 The opendatahub.io Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream_usage_enforcer
+
+import (
+	"context"
+	"encoding/json"
+
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	logutil "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/observability/logging"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling"
+)
+
+const PluginType = "stream-usage-enforcer"
+
+var _ requesthandling.RequestProcessor = &Plugin{}
+
+type Plugin struct {
+	typedName plugin.TypedName
+}
+
+func Factory(name string, _ json.RawMessage, _ plugin.Handle) (plugin.Plugin, error) {
+	return (&Plugin{
+		typedName: plugin.TypedName{Type: PluginType, Name: PluginType},
+	}).WithName(name), nil
+}
+
+func (p *Plugin) WithName(name string) *Plugin {
+	p.typedName.Name = name
+	return p
+}
+
+func (p *Plugin) TypedName() plugin.TypedName { return p.typedName }
+
+func (p *Plugin) ProcessRequest(ctx context.Context, _ *plugin.CycleState, request *requesthandling.InferenceRequest) error {
+	logger := log.FromContext(ctx).V(logutil.VERBOSE)
+
+	stream, ok := request.Body["stream"].(bool)
+	if !ok || !stream {
+		return nil
+	}
+
+	streamOptions, _ := request.Body["stream_options"].(map[string]any)
+	if streamOptions == nil {
+		streamOptions = map[string]any{}
+	}
+
+	if existing, ok := streamOptions["include_usage"].(bool); ok && existing {
+		logger.Info("stream_options.include_usage already set")
+		return nil
+	}
+
+	streamOptions["include_usage"] = true
+	request.SetBodyField("stream_options", streamOptions)
+
+	logger.Info("enforced stream_options.include_usage=true")
+	return nil
+}

pkg/plugins/stream-usage-enforcer/plugin_test.go

@@ -0,0 +1,130 @@
+/*
+Copyright 2026 The opendatahub.io Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package stream_usage_enforcer
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/plugin"
+	"github.com/llm-d/llm-d-inference-payload-processor/pkg/framework/interface/requesthandling"
+)
+
+func TestFactory(t *testing.T) {
+	p, err := Factory("test", nil, nil)
+	require.NoError(t, err)
+	assert.Equal(t, "test", p.TypedName().Name)
+	assert.Equal(t, PluginType, p.TypedName().Type)
+}
+
+func TestProcessRequest_NonStreamingRequest(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["model"] = "gpt-oss-20b"
+	req.Body["messages"] = []any{map[string]any{"role": "user", "content": "Hi"}}
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+	assert.Nil(t, req.Body["stream_options"], "should not add stream_options to non-streaming request")
+	assert.False(t, req.BodyMutated())
+}
+
+func TestProcessRequest_StreamFalse(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["stream"] = false
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+	assert.Nil(t, req.Body["stream_options"])
+	assert.False(t, req.BodyMutated())
+}
+
+func TestProcessRequest_StreamTrueNoStreamOptions(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["model"] = "gpt-oss-20b"
+	req.Body["stream"] = true
+	req.Body["messages"] = []any{map[string]any{"role": "user", "content": "Hello"}}
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+
+	opts, ok := req.Body["stream_options"].(map[string]any)
+	require.True(t, ok, "stream_options should be set")
+	assert.Equal(t, true, opts["include_usage"])
+	assert.True(t, req.BodyMutated())
+}
+
+func TestProcessRequest_StreamTrueIncludeUsageFalse(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["stream"] = true
+	req.Body["stream_options"] = map[string]any{"include_usage": false}
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+
+	opts := req.Body["stream_options"].(map[string]any)
+	assert.Equal(t, true, opts["include_usage"], "should override false to true")
+	assert.True(t, req.BodyMutated())
+}
+
+func TestProcessRequest_StreamTrueIncludeUsageAlreadyTrue(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["stream"] = true
+	req.Body["stream_options"] = map[string]any{"include_usage": true}
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+
+	opts := req.Body["stream_options"].(map[string]any)
+	assert.Equal(t, true, opts["include_usage"])
+	assert.False(t, req.BodyMutated(), "should not mutate when already correct")
+}
+
+func TestProcessRequest_PreservesExistingStreamOptions(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["stream"] = true
+	req.Body["stream_options"] = map[string]any{
+		"continuous_usage_stats": true,
+	}
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+
+	opts := req.Body["stream_options"].(map[string]any)
+	assert.Equal(t, true, opts["include_usage"], "include_usage should be added")
+	assert.Equal(t, true, opts["continuous_usage_stats"], "existing fields should be preserved")
+	assert.True(t, req.BodyMutated())
+}
+
+func TestProcessRequest_StreamFieldAbsent(t *testing.T) {
+	instance, _ := Factory("test", nil, nil)
+	req := requesthandling.NewInferenceRequest()
+	req.Body["model"] = "gpt-oss-20b"
+
+	err := instance.(*Plugin).ProcessRequest(context.Background(), plugin.NewCycleState(), req)
+	require.NoError(t, err)
+	assert.Nil(t, req.Body["stream_options"])
+	assert.False(t, req.BodyMutated())
+}