plugins.go

/*
Copyright 2025 The Kubernetes 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 requestcontrol

import (
	"context"
	"time"

	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
	"github.com/llm-d/llm-d-router/pkg/epp/framework/interface/plugin"
	fwksched "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/scheduling"
)

const (
	PreAdmissionExtensionPoint      = "PreAdmission"
	AdmissionExtensionPoint         = "Admission"
	DataProducerExtensionPoint      = "DataProducer"
	PreRequestExtensionPoint        = "PreRequest"
	ResponseReceivedExtensionPoint  = "ResponseReceived"
	ResponseStreamingExtensionPoint = "ResponseStreaming"
	ResponseCompleteExtensionPoint  = "ResponseComplete"
)

// PreRequest is called by the director after a getting result from scheduling layer and
// before a request is sent to the selected model server.
type PreRequest interface {
	plugin.Plugin
	PreRequest(ctx context.Context, request *fwksched.InferenceRequest, schedulingResult *fwksched.SchedulingResult)
}

// ResponseHeaderProcessor is called by the director after the response headers are successfully received
// which indicates the beginning of the response handling by the model server.
// The given pod argument is the pod that served the request.
type ResponseHeaderProcessor interface {
	plugin.Plugin
	ResponseHeader(ctx context.Context, request *fwksched.InferenceRequest, response *Response, targetEndpoint *datalayer.EndpointMetadata)
}

// ResponseBodyProcessor is the primary hook for processing response data.
// It is called by the director for every data chunk in a streaming response, or exactly once
// for non-streaming responses.
//
// Lifecycle & Termination:
//   - For streams: Invoked multiple times. The final call will have response.EndOfStream set to true.
//   - For non-streaming: Invoked once with response.EndOfStream set to true.
//   - Plugins must treat the call where response.EndOfStream == true as the final lifecycle hook
//     to perform cleanup or final logging.
//
// TODO(https://github.com/kubernetes-sigs/gateway-api-inference-extension/issues/2079):
// Update signature to pass error/termination state. This is a breaking change required for plugins to distinguish
// between success, errors, and disconnects.
type ResponseBodyProcessor interface {
	plugin.Plugin
	ResponseBody(ctx context.Context, request *fwksched.InferenceRequest, response *Response, targetEndpoint *datalayer.EndpointMetadata)
}

// DataProducer is implemented by data producers which produce data from different sources.
// Produce is called by the director before scheduling requests.
type DataProducer interface {
	plugin.ProducerPlugin
	Produce(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) error
}

// TimeoutAwareProducer is an optional interface a DataProducer may implement to
// declare its own execution timeout, overriding the default. A non-positive
// value selects the default.
type TimeoutAwareProducer interface {
	ProduceTimeout() time.Duration
}

// Admitter is called by the director after the data producer and before scheduling.
// When a request has to go through multiple Admitter,
// the request is admitted only if all plugins say that the request should be admitted.
type Admitter interface {
	plugin.Plugin
	// Admit returns the denial reason, wrapped as error if the request is denied.
	// If the request is allowed, it returns nil.
	Admit(ctx context.Context, request *fwksched.InferenceRequest, pods []fwksched.Endpoint) error
}

// PreAdmitter runs after InferenceRequest creation but before admission control.
// It can mutate InferenceRequest fields such as FairnessID and Headers.
type PreAdmitter interface {
	plugin.Plugin
	PreAdmit(ctx context.Context, request *fwksched.InferenceRequest) error
}

// ConditionalDecodeDecider answers the RFC 7240 "Prefer: if-available" gate:
// given the chosen decode endpoint, should the request be rejected with HTTP
// 412 (because the local KV cache does not cover enough of the prompt) or
// forwarded? At most one such plugin is consulted per director.
type ConditionalDecodeDecider interface {
	plugin.Plugin
	ShouldRejectConditionalDecode(ctx context.Context, request *fwksched.InferenceRequest, endpoint fwksched.Endpoint) bool
}