Skip to content

Add configurable prefill port support to prefill-header-handler plugin #329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add configurable prefill port support to prefill-header-handler plugin #329

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e71dcff
Added logic to handle the prefill port
pallavijaini0525 Aug 31, 2025
197299b
Added logic to handle the prefill port - Fixed compilation issues
pallavijaini0525 Aug 31, 2025
d51298c
Edited the comments
pallavijaini0525 Sep 2, 2025
bbfa245
Fixed the lint issues
pallavijaini0525 Oct 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions pkg/plugins/pre-request/pd_prerequest.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ const (
)

type prefillHeaderHandlerParameters struct {
PrefillProfile string `json:"prefillProfile"`
PrefillProfile string `json:"prefillProfile"`
PrefillTargetPort int `json:"prefillTargetPort,omitempty"` // Optional field
}

// compile-time type assertion
Expand All @@ -32,28 +33,31 @@ var _ requestcontrol.PreRequest = &PrefillHeaderHandler{}
// PrefillHeaderHandlerFactory defines the factory function for the PrefillHeaderHandler
func PrefillHeaderHandlerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
parameters := prefillHeaderHandlerParameters{
PrefillProfile: defaultPrefillProfile,
PrefillProfile: defaultPrefillProfile,
PrefillTargetPort: 0,
}
if rawParameters != nil {
if err := json.Unmarshal(rawParameters, &parameters); err != nil {
return nil, fmt.Errorf("failed to parse the parameters of the '%s' pre-request plugin - %w", PrefillHeaderHandlerType, err)
}
}
return NewPrefillHeaderHandler(parameters.PrefillProfile).WithName(name), nil
return NewPrefillHeaderHandler(parameters.PrefillProfile, parameters.PrefillTargetPort).WithName(name), nil
}

// NewPrefillHeaderHandler initializes a new PrefillHeaderHandler and returns its pointer.
func NewPrefillHeaderHandler(prefillProfile string) *PrefillHeaderHandler {
func NewPrefillHeaderHandler(prefillProfile string, prefillTargetPort int) *PrefillHeaderHandler {
return &PrefillHeaderHandler{
typedName: plugins.TypedName{Type: PrefillHeaderHandlerType},
prefillProfile: prefillProfile,
typedName: plugins.TypedName{Type: PrefillHeaderHandlerType},
prefillProfile: prefillProfile,
prefillTargetPort: prefillTargetPort,
}
}

// PrefillHeaderHandler PreRequest plugin
type PrefillHeaderHandler struct {
typedName plugins.TypedName
prefillProfile string
typedName plugins.TypedName
prefillProfile string
prefillTargetPort int // 0 means "not configured"
}

// TypedName returns the typed name of the plugin.
Expand All @@ -78,6 +82,10 @@ func (p *PrefillHeaderHandler) PreRequest(_ context.Context, request *types.LLMR
return // prefill profile failed to run or we chose not to run it, no-op in this case
}

prefillHostPort := net.JoinHostPort(prefillProfileRunResult.TargetPods[0].GetPod().Address, strconv.Itoa(targetPort))
portToUse := targetPort // Default: use existing behavior
if p.prefillTargetPort > 0 { // Only override if configured
portToUse = p.prefillTargetPort
}
prefillHostPort := net.JoinHostPort(prefillProfileRunResult.TargetPods[0].GetPod().Address, strconv.Itoa(portToUse))
request.Headers[prefillPodHeader] = prefillHostPort // in the form of <ip:port>
}