diff --git a/pkg/plugins/pre-request/pd_prerequest.go b/pkg/plugins/pre-request/pd_prerequest.go index 43ecca42..a9e340ce 100644 --- a/pkg/plugins/pre-request/pd_prerequest.go +++ b/pkg/plugins/pre-request/pd_prerequest.go @@ -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 @@ -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, ¶meters); 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. @@ -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 }