feat(edge): OpenAI compatible provider pass-through & test coverage
This commit is contained in:
parent
9cff91f8bf
commit
33f8f5e57f
4 changed files with 927 additions and 3 deletions
|
|
@ -247,6 +247,14 @@ func (s *Server) handleChatCompletionsProviderPool(
|
|||
return rewriteChatCompletionModel(rawBody, target, req)
|
||||
}
|
||||
|
||||
// Normalized branch: chat completions use lenient decode so unknown
|
||||
// fields survive; PrepareRun passes through the already-built Run
|
||||
// without strict validation (SDD S02/S03). The tunnel branch is
|
||||
// unaffected.
|
||||
poolReq.PrepareRun = func(runReq edgeservice.SubmitRunRequest) (edgeservice.SubmitRunRequest, error) {
|
||||
return runReq, nil
|
||||
}
|
||||
|
||||
result, err := s.service.SubmitProviderPool(r.Context(), poolReq)
|
||||
if err != nil {
|
||||
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(req.Model), usageEndpointChatCompletions, responseModePassthrough)
|
||||
|
|
@ -285,6 +293,7 @@ func (s *Server) handleChatCompletionsProviderPool(
|
|||
Run: req,
|
||||
Tunnel: poolReq.Tunnel,
|
||||
PrepareTunnel: poolReq.PrepareTunnel,
|
||||
PrepareRun: poolReq.PrepareRun,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ package openai
|
|||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
|
@ -45,6 +46,19 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
|
|||
writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
|
||||
if dispatch.ProviderPool {
|
||||
// Provider-pool path: the one-shot SubmitProviderPool selects a candidate
|
||||
// and dispatches either tunnel passthrough (OpenAI-compatible) or
|
||||
// normalized RunEvent (Ollama/CLI/native) based on the candidate's
|
||||
// executionPath. Raw body is preserved for the tunnel branch; the
|
||||
// normalized branch performs strict decode + prompt build below.
|
||||
s.handleResponsesProviderPool(r, rawBody, env, dispatch, w)
|
||||
return
|
||||
}
|
||||
|
||||
// Non-provider routes: tunnel or normalized. Tunnel routes relay the raw
|
||||
// body verbatim; normalized routes require strict decode + prompt build.
|
||||
if routeUsesProviderTunnel(dispatch) {
|
||||
runMeta, workspace, err := parseOpenAIMetadata(env.Metadata)
|
||||
if err != nil {
|
||||
|
|
@ -69,8 +83,8 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
|
|||
return
|
||||
}
|
||||
|
||||
// Non-provider routes keep the normalized RunEvent path: strict decode,
|
||||
// stream/background rejection, prompt build, and SubmitRun.
|
||||
// Normalized RunEvent path: strict decode, stream/background rejection,
|
||||
// prompt build, and SubmitRun.
|
||||
var req responsesRequest
|
||||
if err := decodeResponsesRequest(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
|
|
@ -190,6 +204,269 @@ func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
|
|||
s.completeResponse(w, r, req, handle, outputPolicy)
|
||||
}
|
||||
|
||||
// handleResponsesProviderPool dispatches a provider-pool /v1/responses request
|
||||
// through the one-shot SubmitProviderPool surface. After a single queue
|
||||
// admission, the selected candidate's executionPath determines whether the
|
||||
// response is delivered via tunnel passthrough (OpenAI-compatible providers)
|
||||
// or the normalized RunEvent path (Ollama/CLI/native). Provider auth is
|
||||
// validated via providerTunnelAuthHeaders for tunnel-path only (SDD D01).
|
||||
// Strict-output policy only applies to normalized dispatch, not to raw tunnel
|
||||
// passthrough (SDD D02).
|
||||
//
|
||||
// Selection-first flow: the provider is selected by SubmitProviderPool before
|
||||
// any path-specific decode, validation, or stream gating. The tunnel branch
|
||||
// only rewrites the model field and relays the raw body (SDD S04). The
|
||||
// normalized branch performs strict decode + stream/background rejection +
|
||||
// prompt build, but only after the candidate is known. This prevents
|
||||
// tunnel-preferring unknown or streaming requests from bypassing the selected
|
||||
// provider's executionPath rule (SDD S02/S03).
|
||||
func (s *Server) handleResponsesProviderPool(
|
||||
r *http.Request,
|
||||
rawBody []byte,
|
||||
env responsesEnvelope,
|
||||
dispatch routeDispatch,
|
||||
w http.ResponseWriter,
|
||||
) {
|
||||
// Build a populated base Run for the normalized branch and a Tunnel for
|
||||
// the tunnel branch. Both are sent to SubmitProviderPool so the one-shot
|
||||
// admission selects the provider before any path-specific behavior.
|
||||
runMeta, workspace, err := parseOpenAIMetadata(env.Metadata)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
for k, v := range principalMetadata(r.Context()) {
|
||||
runMeta[k] = v
|
||||
}
|
||||
if err := validateWorkspaceForRoute(dispatch, workspace); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Compute estimate and context class from the raw body so both tunnel and
|
||||
// normalized branches carry the same observability / long-context
|
||||
// information the direct tunnel path carries (SDD S02/S03).
|
||||
estimate := estimateInputTokens(string(rawBody), runMeta, nil, nil)
|
||||
contextClass := classifyContext(estimate, s.longContextThreshold())
|
||||
runMeta["openai_model"] = env.Model
|
||||
runMeta["openai_stream"] = strconv.FormatBool(env.Stream)
|
||||
runMeta["estimated_input_tokens"] = strconv.Itoa(estimate)
|
||||
runMeta["context_class"] = contextClass
|
||||
|
||||
// Base Run carries the model-group key that SubmitProviderPool requires
|
||||
// for candidate resolution. Metadata, estimate, and context class are
|
||||
// preserved by PrepareRun on the normalized branch; the tunnel branch
|
||||
// inherits them from the base so tunnel dispatch keeps observability.
|
||||
baseRun := edgeservice.SubmitRunRequest{
|
||||
NodeRef: dispatch.NodeRef,
|
||||
ModelGroupKey: strings.TrimSpace(env.Model),
|
||||
SessionID: dispatch.SessionID,
|
||||
Workspace: workspace,
|
||||
Metadata: runMeta,
|
||||
EstimatedInputTokens: estimate,
|
||||
ContextClass: contextClass,
|
||||
ProviderPool: true,
|
||||
}
|
||||
|
||||
baseTunnel := edgeservice.SubmitProviderTunnelRequest{
|
||||
Metadata: runMeta,
|
||||
EstimatedInputTokens: estimate,
|
||||
ContextClass: contextClass,
|
||||
ModelGroupKey: strings.TrimSpace(env.Model),
|
||||
SessionID: dispatch.SessionID,
|
||||
Method: http.MethodPost,
|
||||
Path: "/v1/responses",
|
||||
Stream: env.Stream,
|
||||
TimeoutSec: dispatch.TimeoutSec,
|
||||
MaxQueue: dispatch.MaxQueue,
|
||||
QueueTimeoutMS: dispatch.QueueTimeoutMS,
|
||||
ProviderPool: true,
|
||||
}
|
||||
|
||||
poolReq := edgeservice.ProviderPoolDispatchRequest{
|
||||
Run: baseRun,
|
||||
Tunnel: baseTunnel,
|
||||
}
|
||||
|
||||
// Pre-dispatch provider auth header injection. Runs inside SubmitProviderPool
|
||||
// BEFORE buildProviderTunnelRequest and the Node Send step, so the auth
|
||||
// header lands on the actual wire request. On failure the slot is released
|
||||
// and no tunnel request is sent (SDD S03).
|
||||
poolReq.PrepareTunnel = func(tunnelReq edgeservice.SubmitProviderTunnelRequest) (edgeservice.SubmitProviderTunnelRequest, error) {
|
||||
headers, err := s.providerTunnelAuthHeaders(r)
|
||||
if err != nil {
|
||||
return tunnelReq, err
|
||||
}
|
||||
tunnelReq.Headers = headers
|
||||
return tunnelReq, nil
|
||||
}
|
||||
|
||||
// Tunnel branch: rewrite only the model field, preserve all other fields
|
||||
// (tools, max_output_tokens, custom fields). Stream/background gating is
|
||||
// skipped: the provider itself enforces those constraints (SDD S04).
|
||||
poolReq.Tunnel.BuildBody = func(target string) ([]byte, error) {
|
||||
return rewriteResponsesModel(rawBody, target)
|
||||
}
|
||||
|
||||
// Prepared values captured from the PrepareRun hook on the normalized path.
|
||||
// PrepareRun runs strict decode + prompt build + output policy; the
|
||||
// normalized branch must relay these exact values to completeResponse so
|
||||
// strict-output XML wrapping survives the provider-pool path.
|
||||
var (
|
||||
preparedResponsesReq responsesRequest
|
||||
preparedOutputPolicy strictOutputPolicy
|
||||
prepareRunSucceeded bool
|
||||
)
|
||||
|
||||
// Normalized branch: strict decode, stream/background rejection, prompt
|
||||
// build, and generation policy. Only runs AFTER the candidate is selected
|
||||
// by SubmitProviderPool (i.e. only when executionPath is normalized).
|
||||
// Validation errors are wrapped with errProviderRequestValidation so
|
||||
// SubmitProviderPool can relay them as HTTP 400.
|
||||
poolReq.PrepareRun = func(runReq edgeservice.SubmitRunRequest) (edgeservice.SubmitRunRequest, error) {
|
||||
req := responsesRequest{Model: env.Model}
|
||||
if err := decodeResponsesRequest(json.NewDecoder(bytes.NewReader(rawBody)), &req); err != nil {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: invalid /v1/responses request format: %s", errProviderRequestValidation, err.Error())
|
||||
}
|
||||
if req.Stream {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: streaming is not supported for /v1/responses", errProviderRequestValidation)
|
||||
}
|
||||
if req.Background {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: background is not supported for /v1/responses", errProviderRequestValidation)
|
||||
}
|
||||
|
||||
inputStr, parseInputErr := parseResponsesInput(req.Input)
|
||||
if parseInputErr != nil {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: %s", errProviderRequestValidation, parseInputErr.Error())
|
||||
}
|
||||
|
||||
prompt := buildResponsesPrompt(req.Instructions, inputStr)
|
||||
outputPolicy := s.resolveOutputPolicy(prompt)
|
||||
if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
|
||||
prompt = instruction + "\n" + prompt
|
||||
}
|
||||
|
||||
preparedMeta, _, metaErr := parseOpenAIMetadata(req.Metadata)
|
||||
if metaErr != nil {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: %s", errProviderRequestValidation, metaErr.Error())
|
||||
}
|
||||
for k, v := range principalMetadata(r.Context()) {
|
||||
preparedMeta[k] = v
|
||||
}
|
||||
|
||||
if workspaceErr := validateWorkspaceForRoute(dispatch, workspace); workspaceErr != nil {
|
||||
return edgeservice.SubmitRunRequest{}, fmt.Errorf("%w: %s", errProviderRequestValidation, workspaceErr.Error())
|
||||
}
|
||||
|
||||
var defaultThinkingTokenBudget int
|
||||
if catalogEntry := s.findProviderPoolEntry(req.Model); catalogEntry != nil {
|
||||
applyModelCatalogGenerationPolicyToResponses(&req, *catalogEntry)
|
||||
if catalogEntry.DefaultThinkingTokenBudget > 0 {
|
||||
defaultThinkingTokenBudget = catalogEntry.DefaultThinkingTokenBudget
|
||||
}
|
||||
}
|
||||
|
||||
preparedMeta["openai_model"] = req.Model
|
||||
preparedMeta["openai_stream"] = fmt.Sprintf("%t", req.Stream)
|
||||
preparedMeta["strict_output"] = fmt.Sprintf("%t", outputPolicy.Strict)
|
||||
|
||||
input := map[string]any{"prompt": prompt}
|
||||
if defaultThinkingTokenBudget > 0 {
|
||||
input["think"] = true
|
||||
input["thinking_token_budget"] = defaultThinkingTokenBudget
|
||||
} else if outputPolicy.Strict {
|
||||
input["think"] = false
|
||||
}
|
||||
if options := req.providerOptions(); len(options) > 0 {
|
||||
input["options"] = options
|
||||
}
|
||||
|
||||
estimate := estimateInputTokens(prompt, preparedMeta, nil, nil)
|
||||
contextClass := classifyContext(estimate, s.longContextThreshold())
|
||||
preparedMeta["estimated_input_tokens"] = strconv.Itoa(estimate)
|
||||
preparedMeta["context_class"] = contextClass
|
||||
|
||||
runReq.Prompt = prompt
|
||||
runReq.Input = input
|
||||
runReq.Metadata = preparedMeta
|
||||
runReq.EstimatedInputTokens = estimate
|
||||
runReq.ContextClass = contextClass
|
||||
runReq.TimeoutSec = dispatch.TimeoutSec
|
||||
runReq.MaxQueue = dispatch.MaxQueue
|
||||
runReq.QueueTimeoutMS = dispatch.QueueTimeoutMS
|
||||
|
||||
preparedResponsesReq = req
|
||||
preparedOutputPolicy = outputPolicy
|
||||
prepareRunSucceeded = true
|
||||
return runReq, nil
|
||||
}
|
||||
|
||||
result, err := s.service.SubmitProviderPool(r.Context(), poolReq)
|
||||
if err != nil {
|
||||
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(env.Model), usageEndpointResponses, responseModePassthrough)
|
||||
emitUsageMetrics(metricLabels, usageStatusForError(err), usageObservation{})
|
||||
// Provider auth failure is a client request error (400), not a
|
||||
// backend dispatch error. The auth check runs inside
|
||||
// SubmitProviderPool (PrepareTunnel) before any tunnel request is sent.
|
||||
// Normalized validation failures (e.g. strict decode, stream) are
|
||||
// also client errors: map to 400.
|
||||
if errors.Is(err, errProviderAuthRequired) {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", "provider auth token is required")
|
||||
return
|
||||
}
|
||||
if isValidationError(err) {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
s.logger.Info("openai responses provider-pool dispatch",
|
||||
zap.String("run_id", result.DispatchInfo.RunID),
|
||||
zap.String("node_id", result.DispatchInfo.NodeID),
|
||||
zap.String("provider_id", result.DispatchInfo.ProviderID),
|
||||
zap.String("provider_type", result.DispatchInfo.ProviderType),
|
||||
zap.String("execution_path", result.DispatchInfo.ExecutionPath),
|
||||
zap.String("model_group", result.DispatchInfo.ModelGroupKey),
|
||||
zap.String("adapter", result.DispatchInfo.Adapter),
|
||||
zap.String("target", result.DispatchInfo.Target),
|
||||
zap.String("path", string(result.Path)),
|
||||
zap.Int("estimated_input_tokens", result.DispatchInfo.EstimatedInputTokens),
|
||||
zap.String("context_class", result.DispatchInfo.ContextClass),
|
||||
zap.String("queue_reason", result.DispatchInfo.QueueReason),
|
||||
)
|
||||
|
||||
switch result.Path {
|
||||
case edgeservice.ProviderPoolPathTunnel:
|
||||
// Tunnel path: provider auth was already validated and injected via
|
||||
// PrepareTunnel before dispatch; on failure SubmitProviderPool returns
|
||||
// an error and no tunnel handle exists. Provider bytes are relayed as
|
||||
// pure passthrough; caller metadata never selects a sideband surface.
|
||||
metricLabels := s.usageLabelsFor(r.Context(), strings.TrimSpace(env.Model), usageEndpointResponses, responseModePassthrough)
|
||||
s.writeProviderTunnelResponse(w, r, result.Tunnel, env.Stream, env.Model, metricLabels)
|
||||
|
||||
case edgeservice.ProviderPoolPathNormalized:
|
||||
// Normalized path: no auth required, collect from RunEvent stream.
|
||||
handle := result.Run
|
||||
if handle == nil {
|
||||
writeError(w, http.StatusInternalServerError, "run_error", "provider-pool selection returned normalized path but no run result")
|
||||
return
|
||||
}
|
||||
// PrepareRun is only invoked on the normalized path. A normalized
|
||||
// result without a successful PrepareRun is an internal mismatch and
|
||||
// must surface as run_error (SDD D02).
|
||||
if !prepareRunSucceeded {
|
||||
writeError(w, http.StatusInternalServerError, "run_error", "provider-pool normalized result received without a successful PrepareRun")
|
||||
return
|
||||
}
|
||||
// Relay the prepared request and output policy from PrepareRun to
|
||||
// completeResponse so strict-output XML wrapping survives the
|
||||
// normalized provider-pool path.
|
||||
s.completeResponse(w, r, preparedResponsesReq, handle, preparedOutputPolicy)
|
||||
}
|
||||
}
|
||||
|
||||
func decodeResponsesRequest(dec *json.Decoder, req *responsesRequest) error {
|
||||
var raw map[string]json.RawMessage
|
||||
if err := dec.Decode(&raw); err != nil {
|
||||
|
|
@ -441,6 +718,22 @@ func parseOpenAIMetadata(raw json.RawMessage) (map[string]string, string, error)
|
|||
return flat, workspace, nil
|
||||
}
|
||||
|
||||
// errProviderRequestValidation is a sentinel for client request validation
|
||||
// failures that occur inside the normalized provider-pool PrepareRun hook
|
||||
// (strict decode, stream/background rejection, input parsing). These errors
|
||||
// are reported as HTTP 400 invalid_request_error, distinct from backend
|
||||
// dispatch errors that return 502.
|
||||
var errProviderRequestValidation = errors.New("provider_request_validation")
|
||||
|
||||
// isValidationError reports whether err originated from a PrepareRun validation
|
||||
// failure in the normalized provider-pool path.
|
||||
func isValidationError(err error) bool {
|
||||
return errors.Is(err, errProviderRequestValidation)
|
||||
}
|
||||
|
||||
// metadataStringValue extracts a single string value from a json.RawMessage
|
||||
// for the metadata field. It rejects non-string values and enforces the
|
||||
// 512-character length cap.
|
||||
func metadataStringValue(key string, raw json.RawMessage) (string, error) {
|
||||
var value string
|
||||
if err := json.Unmarshal(raw, &value); err != nil {
|
||||
|
|
|
|||
|
|
@ -71,6 +71,18 @@ type fakeRunService struct {
|
|||
// calls. When non-empty, results are consumed in order; when empty, the
|
||||
// poolDispatchPath / poolSubmitErr fields are used instead.
|
||||
poolSubmitResults []edgeservice.ProviderPoolDispatchResult
|
||||
// poolSubmitCount tracks how many times SubmitProviderPool was called
|
||||
// (including from the chat handler). Used by tests to assert dispatch
|
||||
// surface fidelity.
|
||||
poolSubmitCount int
|
||||
// poolLastRun records the last SubmitRunRequest passed to SubmitProviderPool
|
||||
// (pre-PrepareRun). Tests assert ModelGroupKey is non-empty to catch the
|
||||
// "empty Run" regression.
|
||||
poolLastRun edgeservice.SubmitRunRequest
|
||||
// poolPrepareRunCalled tracks whether the optional PrepareRun hook was
|
||||
// invoked by SubmitProviderPool on the normalized path. Only meaningful
|
||||
// when poolDispatchPath == "normalized".
|
||||
poolPrepareRunCalled bool
|
||||
}
|
||||
|
||||
type fakeTunnelHandle struct {
|
||||
|
|
@ -158,7 +170,15 @@ func (s *fakeRunService) SubmitProviderTunnel(_ context.Context, req edgeservice
|
|||
// SubmitProviderPool is a minimal test double for the provider-pool one-shot
|
||||
// dispatch. It records the selected path (tunnel vs normalized) and returns a
|
||||
// stub result so the OpenAI handler can exercise the selection-first flow.
|
||||
// It also tracks whether the optional PrepareRun hook was invoked, so tests
|
||||
// can assert the handler only calls PrepareRun on the normalized path.
|
||||
func (s *fakeRunService) SubmitProviderPool(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
|
||||
s.submitMu.Lock()
|
||||
s.poolSubmitCount++
|
||||
s.poolLastRun = req.Run
|
||||
s.poolPrepareRunCalled = false
|
||||
s.submitMu.Unlock()
|
||||
|
||||
// Custom results for successive calls (e.g. tool-validation retry tests).
|
||||
s.submitMu.Lock()
|
||||
if len(s.poolSubmitResults) > 0 {
|
||||
|
|
@ -254,6 +274,20 @@ func (s *fakeRunService) SubmitProviderPool(_ context.Context, req edgeservice.P
|
|||
frames = relayed
|
||||
} else if frames == nil {
|
||||
frames = staticProviderTunnelFrames(`{"ok":true}`)
|
||||
// Record the built tunnel body even when using a static frame source.
|
||||
body := req.Tunnel.Body
|
||||
if req.Tunnel.BuildBody != nil {
|
||||
target := req.Tunnel.Target
|
||||
if s.tunnelServedTarget != "" {
|
||||
target = s.tunnelServedTarget
|
||||
}
|
||||
if built, buildErr := req.Tunnel.BuildBody(target); buildErr == nil {
|
||||
body = built
|
||||
}
|
||||
}
|
||||
s.submitMu.Lock()
|
||||
s.tunnelBodies = append(s.tunnelBodies, body)
|
||||
s.submitMu.Unlock()
|
||||
} else {
|
||||
body := req.Tunnel.Body
|
||||
if req.Tunnel.BuildBody != nil {
|
||||
|
|
@ -295,9 +329,22 @@ func (s *fakeRunService) SubmitProviderPool(_ context.Context, req edgeservice.P
|
|||
}
|
||||
|
||||
// Normalized path (mimics Ollama/CLI behavior).
|
||||
// Invoke the optional PrepareRun hook so tests can assert it only runs
|
||||
// on the normalized path (mirrors real SubmitProviderPool behavior).
|
||||
if req.PrepareRun != nil {
|
||||
s.submitMu.Lock()
|
||||
s.poolPrepareRunCalled = true
|
||||
s.submitMu.Unlock()
|
||||
prepared, prepErr := req.PrepareRun(req.Run)
|
||||
if prepErr != nil {
|
||||
return nil, prepErr
|
||||
}
|
||||
req.Run = prepared
|
||||
}
|
||||
s.submitMu.Lock()
|
||||
s.reqs = append(s.reqs, req.Run)
|
||||
s.submitMu.Unlock()
|
||||
|
||||
runChan := s.poolRunFrames
|
||||
if runChan == nil {
|
||||
runChan = make(chan *iop.RunEvent, 1)
|
||||
|
|
@ -376,6 +423,18 @@ func (s *fakeRunService) buildPoolResult(req edgeservice.ProviderPoolDispatchReq
|
|||
if disp.NodeID == "" {
|
||||
disp.NodeID = "node-pool"
|
||||
}
|
||||
// Invoke the optional PrepareRun hook on the normalized path (mirrors
|
||||
// real SubmitProviderPool behavior).
|
||||
if req.PrepareRun != nil {
|
||||
prepared, prepErr := req.PrepareRun(req.Run)
|
||||
if prepErr != nil {
|
||||
return nil, prepErr
|
||||
}
|
||||
req.Run = prepared
|
||||
}
|
||||
s.submitMu.Lock()
|
||||
s.poolPrepareRunCalled = true
|
||||
s.submitMu.Unlock()
|
||||
s.submitMu.Lock()
|
||||
s.reqs = append(s.reqs, req.Run)
|
||||
s.submitMu.Unlock()
|
||||
|
|
@ -423,6 +482,32 @@ func (s *fakeRunService) tunnelHandleSnapshot() tunnelHandleSnapshot {
|
|||
}
|
||||
}
|
||||
|
||||
// poolSubmitCountSnapshot returns the number of times SubmitProviderPool was
|
||||
// called. Tests assert exactly-once dispatch from the handler.
|
||||
func (s *fakeRunService) poolSubmitCountSnapshot() int {
|
||||
s.submitMu.Lock()
|
||||
defer s.submitMu.Unlock()
|
||||
return s.poolSubmitCount
|
||||
}
|
||||
|
||||
// poolLastRunSnapshot returns a copy of the last SubmitRunRequest passed to
|
||||
// SubmitProviderPool. Tests assert ModelGroupKey is non-empty to catch the
|
||||
// "empty Run" regression.
|
||||
func (s *fakeRunService) poolLastRunSnapshot() edgeservice.SubmitRunRequest {
|
||||
s.submitMu.Lock()
|
||||
defer s.submitMu.Unlock()
|
||||
return s.poolLastRun
|
||||
}
|
||||
|
||||
// poolPrepareRunCalledSnapshot returns whether the PrepareRun hook was
|
||||
// invoked by the last SubmitProviderPool call. Only meaningful when the last
|
||||
// call was on the normalized path.
|
||||
func (s *fakeRunService) poolPrepareRunCalledSnapshot() bool {
|
||||
s.submitMu.Lock()
|
||||
defer s.submitMu.Unlock()
|
||||
return s.poolPrepareRunCalled
|
||||
}
|
||||
|
||||
func (s *fakeRunService) tunnelReqsSnapshot() []edgeservice.SubmitProviderTunnelRequest {
|
||||
s.submitMu.Lock()
|
||||
reqs := make([]edgeservice.SubmitProviderTunnelRequest, len(s.tunnelReqs))
|
||||
|
|
@ -5676,6 +5761,206 @@ func responsesProviderTunnelServer(frames chan *iop.ProviderTunnelFrame, servedT
|
|||
return srv, fake
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolTunnelSelectionUsesPassthrough verifies that when a
|
||||
// provider-pool catalog match is resolved and the selected provider is an
|
||||
// OpenAI-compatible type (tunnel executionPath), the Responses handler
|
||||
// dispatches via tunnel passthrough with the raw body preserved (model rewrite
|
||||
// only). This mirrors TestChatCompletionsProviderPoolTunnelSelectionUsesPassthrough
|
||||
// for /v1/responses.
|
||||
func TestResponsesProviderPoolTunnelSelectionUsesPassthrough(t *testing.T) {
|
||||
frames := make(chan *iop.ProviderTunnelFrame, 3)
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(`{"ok":true}`)}
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
||||
close(frames)
|
||||
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
||||
tunnelFrames: frames,
|
||||
}
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{ID: "vllm-resp-model", Providers: map[string]string{"prov-vllm": "Qwen3-35B"}},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"vllm-resp-model",
|
||||
"input":"hello world"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
// Tunnel passthrough must use SubmitProviderPool tunnel path.
|
||||
tunnelReqs := fake.tunnelReqsSnapshot()
|
||||
reqs := fake.reqsSnapshot()
|
||||
if len(tunnelReqs) != 1 {
|
||||
t.Fatalf("expected 1 tunnel dispatch for vLLM path, got %d", len(tunnelReqs))
|
||||
}
|
||||
if len(reqs) != 0 {
|
||||
t.Fatalf("expected 0 normalized SubmitRun for vLLM tunnel path, got %d", len(reqs))
|
||||
}
|
||||
if !tunnelReqs[0].ProviderPool {
|
||||
t.Error("tunnel dispatch must have ProviderPool=true")
|
||||
}
|
||||
if tunnelReqs[0].Path != "/v1/responses" {
|
||||
t.Fatalf("tunnel path: got %q want /v1/responses", tunnelReqs[0].Path)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolOllamaSelectionUsesNormalizedRun verifies that when a
|
||||
// provider-pool catalog match is resolved and the selected provider is an
|
||||
// Ollama type (normalized executionPath), the Responses handler dispatches via
|
||||
// normalized RunEvent path with the decoded prompt/input, not via tunnel. This
|
||||
// mirrors TestChatCompletionsProviderPoolOllamaSelectionUsesNormalizedRun for
|
||||
// /v1/responses.
|
||||
func TestResponsesProviderPoolOllamaSelectionUsesNormalizedRun(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
|
||||
}
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{ID: "ollama-resp-model", Providers: map[string]string{"prov-ollama": "qwen3.6:35b"}},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"ollama-resp-model",
|
||||
"input":"hello world"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
// Normalized Ollama path must use SubmitProviderPool normalized path.
|
||||
reqs := fake.reqsSnapshot()
|
||||
tunnelReqs := fake.tunnelReqsSnapshot()
|
||||
if len(reqs) != 1 {
|
||||
t.Fatalf("expected 1 normalized SubmitRun for Ollama path, got %d", len(reqs))
|
||||
}
|
||||
if len(tunnelReqs) != 0 {
|
||||
t.Fatalf("expected 0 tunnel dispatch for Ollama normalized path, got %d", len(tunnelReqs))
|
||||
}
|
||||
if !reqs[0].ProviderPool {
|
||||
t.Error("normalized dispatch must have ProviderPool=true")
|
||||
}
|
||||
if reqs[0].ModelGroupKey != "ollama-resp-model" {
|
||||
t.Fatalf("expected model group key %q, got %q", "ollama-resp-model", reqs[0].ModelGroupKey)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesMixedProviderPoolTunnelSelection verifies S05: a mixed
|
||||
// provider-pool model group that selects an OpenAI-compatible (vLLM) provider
|
||||
// dispatches via tunnel passthrough for /v1/responses.
|
||||
func TestResponsesMixedProviderPoolTunnelSelection(t *testing.T) {
|
||||
tunnelFrames := make(chan *iop.ProviderTunnelFrame, 3)
|
||||
tunnelFrames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
|
||||
tunnelFrames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(`{"id":"resp-mixed","object":"response"}`)}
|
||||
tunnelFrames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
||||
close(tunnelFrames)
|
||||
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
||||
tunnelFrames: tunnelFrames,
|
||||
}
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{
|
||||
ID: "mixed-resp-model",
|
||||
Providers: map[string]string{
|
||||
"prov-vllm": "served-vllm",
|
||||
"prov-ollama": "served-ollama",
|
||||
},
|
||||
},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"mixed-resp-model",
|
||||
"input":"hello mixed"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
tunnelReqs := fake.tunnelReqsSnapshot()
|
||||
if len(tunnelReqs) != 1 {
|
||||
t.Fatalf("expected 1 tunnel dispatch, got %d", len(tunnelReqs))
|
||||
}
|
||||
if !tunnelReqs[0].ProviderPool {
|
||||
t.Error("tunnel request must have ProviderPool=true")
|
||||
}
|
||||
if tunnelReqs[0].Path != "/v1/responses" {
|
||||
t.Fatalf("tunnel path: got %q want /v1/responses", tunnelReqs[0].Path)
|
||||
}
|
||||
|
||||
runReqs := fake.reqsSnapshot()
|
||||
if len(runReqs) != 0 {
|
||||
t.Errorf("expected 0 normalized SubmitRun calls, got %d", len(runReqs))
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesMixedProviderPoolOllamaSelection verifies S05: when the mixed
|
||||
// model group selects an Ollama provider, the request is dispatched over the
|
||||
// normalized path (SubmitRun), and no tunnel dispatch is issued for /v1/responses.
|
||||
func TestResponsesMixedProviderPoolOllamaSelection(t *testing.T) {
|
||||
events := bufferedRunEvents(
|
||||
&iop.RunEvent{Type: "delta", Delta: "ollama-resp"},
|
||||
&iop.RunEvent{Type: "complete"},
|
||||
)
|
||||
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
|
||||
events: events,
|
||||
}
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{
|
||||
ID: "mixed-resp-model",
|
||||
Providers: map[string]string{
|
||||
"prov-vllm": "served-vllm",
|
||||
"prov-ollama": "served-ollama",
|
||||
},
|
||||
},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"mixed-resp-model",
|
||||
"input":"hello mixed"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
tunnelReqs := fake.tunnelReqsSnapshot()
|
||||
if len(tunnelReqs) != 0 {
|
||||
t.Fatalf("expected 0 tunnel dispatches for Ollama selection, got %d", len(tunnelReqs))
|
||||
}
|
||||
|
||||
runReqs := fake.reqsSnapshot()
|
||||
if len(runReqs) != 1 {
|
||||
t.Fatalf("expected 1 normalized SubmitRun, got %d", len(runReqs))
|
||||
}
|
||||
if runReqs[0].ModelGroupKey != "mixed-resp-model" {
|
||||
t.Errorf("normalized ModelGroupKey: got %q, want mixed-resp-model", runReqs[0].ModelGroupKey)
|
||||
}
|
||||
if !runReqs[0].ProviderPool {
|
||||
t.Error("normalized run must have ProviderPool=true")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolDispatch verifies that /v1/responses sends
|
||||
// provider-pool models through the raw provider tunnel (POST /v1/responses)
|
||||
// instead of the normalized RunEvent path.
|
||||
|
|
@ -7440,3 +7725,321 @@ func TestProviderPoolStandardResponseNoExtensionFields(t *testing.T) {
|
|||
t.Error("expected provider tunnel dispatch")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolTunnelUnknownFieldsUsesPoolSelection asserts that
|
||||
// when the selected provider is tunnel-passthrough (OpenAI-compatible),
|
||||
// unknown fields (tools, parallel_tool_calls, store, custom field) are
|
||||
// preserved verbatim in the tunnel body, and that the Run sent to
|
||||
// SubmitProviderPool carries a non-empty ModelGroupKey.
|
||||
func TestResponsesProviderPoolTunnelUnknownFieldsUsesPoolSelection(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
||||
}
|
||||
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog([]config.ModelCatalogEntry{
|
||||
{ID: "pool-model", Providers: map[string]string{"prov-tunnel": "served-tunnel"}},
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"pool-model",
|
||||
"input":"hello world",
|
||||
"tools":[{"type":"function","function":{"name":"lookup"}}],
|
||||
"parallel_tool_calls":true,
|
||||
"store":false,
|
||||
"custom_field":"keep-me"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// SubmitProviderPool must have been called with a populated Run.
|
||||
lastRun := fake.poolLastRunSnapshot()
|
||||
if lastRun.ModelGroupKey == "" {
|
||||
t.Error("SubmitProviderPool was called with empty Run.ModelGroupKey")
|
||||
}
|
||||
|
||||
// The tunnel body must preserve unknown fields verbatim.
|
||||
if len(fake.tunnelBodiesSnapshot()) != 1 {
|
||||
t.Fatalf("expected exactly one tunnel body, got %d", len(fake.tunnelBodiesSnapshot()))
|
||||
}
|
||||
body := fake.tunnelBodiesSnapshot()[0]
|
||||
if !strings.Contains(string(body), "\"tools\"") {
|
||||
t.Error("tunnel body must preserve tools field")
|
||||
}
|
||||
if !strings.Contains(string(body), "\"parallel_tool_calls\"") {
|
||||
t.Error("tunnel body must preserve parallel_tool_calls field")
|
||||
}
|
||||
if !strings.Contains(string(body), "\"store\"") {
|
||||
t.Error("tunnel body must preserve store field")
|
||||
}
|
||||
if !strings.Contains(string(body), "\"custom_field\"") {
|
||||
t.Error("tunnel body must preserve custom_field")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolOllamaSelectionRejectsUnknownFieldsWithoutTunnel
|
||||
// asserts that when the selected provider is normalized (Ollama/CLI), an
|
||||
// unknown-field request returns HTTP 400 and no tunnel or run dispatch occurs.
|
||||
func TestResponsesProviderPoolOllamaSelectionRejectsUnknownFieldsWithoutTunnel(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
|
||||
}
|
||||
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog([]config.ModelCatalogEntry{
|
||||
{ID: "ollama-model", Providers: map[string]string{"prov-ollama": "served-ollama"}},
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"ollama-model",
|
||||
"input":"hello",
|
||||
"unknown_field":"reject-me"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status: got %d body=%s, want 400", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// No tunnel or run dispatch should have occurred.
|
||||
if len(fake.tunnelReqsSnapshot()) != 0 {
|
||||
t.Error("normalized path must not dispatch tunnel for unknown fields")
|
||||
}
|
||||
if len(fake.reqsSnapshot()) != 0 {
|
||||
t.Error("normalized path must not dispatch run for unknown fields")
|
||||
}
|
||||
|
||||
// Verify pool dispatch was made.
|
||||
if fake.poolSubmitCountSnapshot() != 1 {
|
||||
t.Fatalf("pool submit count: got %d, want 1", fake.poolSubmitCountSnapshot())
|
||||
}
|
||||
|
||||
// Verify PrepareRun was called on the normalized path.
|
||||
prepareRunCalled := fake.poolPrepareRunCalledSnapshot()
|
||||
if !prepareRunCalled {
|
||||
t.Errorf("PrepareRun must be called on normalized path (poolSubmitCount=%d, poolLastRun.ModelGroupKey=%q)", fake.poolSubmitCountSnapshot(), fake.poolLastRunSnapshot().ModelGroupKey)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolOllamaSelectionRejectsStreamingWithoutTunnel asserts
|
||||
// that when the selected provider is normalized, stream:true returns HTTP 400
|
||||
// and no tunnel dispatch occurs.
|
||||
func TestResponsesProviderPoolOllamaSelectionRejectsStreamingWithoutTunnel(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
|
||||
}
|
||||
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog([]config.ModelCatalogEntry{
|
||||
{ID: "ollama-model", Providers: map[string]string{"prov-ollama": "served-ollama"}},
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"ollama-model",
|
||||
"input":"hello",
|
||||
"stream":true
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status: got %d body=%s, want 400", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// No tunnel dispatch for normalized path with stream:true.
|
||||
if len(fake.tunnelReqsSnapshot()) != 0 {
|
||||
t.Error("normalized path must not dispatch tunnel for stream:true")
|
||||
}
|
||||
if len(fake.reqsSnapshot()) != 0 {
|
||||
t.Error("normalized path must not dispatch run for stream:true")
|
||||
}
|
||||
|
||||
if !fake.poolPrepareRunCalledSnapshot() {
|
||||
t.Error("PrepareRun must be called on normalized path")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolTunnelStreamingUsesPoolSelection asserts that when
|
||||
// the selected provider is tunnel-passthrough, stream:true is relayed through
|
||||
// SubmitProviderPool with the raw SSE body preserved.
|
||||
func TestResponsesProviderPoolTunnelStreamingUsesPoolSelection(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
||||
tunnelFrames: staticProviderTunnelFrames(`{"id":"resp-1","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"hello"}]}],"stream":true}\n`),
|
||||
}
|
||||
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog([]config.ModelCatalogEntry{
|
||||
{ID: "pool-model", Providers: map[string]string{"prov-tunnel": "served-tunnel"}},
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"pool-model",
|
||||
"input":"hello",
|
||||
"stream":true
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// SubmitProviderPool must have been called.
|
||||
if fake.poolSubmitCountSnapshot() != 1 {
|
||||
t.Fatalf("SubmitProviderPool call count: got %d, want 1", fake.poolSubmitCountSnapshot())
|
||||
}
|
||||
|
||||
// The tunnel body must preserve the raw SSE payload.
|
||||
if len(fake.tunnelBodiesSnapshot()) != 1 {
|
||||
t.Fatalf("expected one tunnel body, got %d", len(fake.tunnelBodiesSnapshot()))
|
||||
}
|
||||
body := fake.tunnelBodiesSnapshot()[0]
|
||||
if !strings.Contains(string(body), "\"stream\":true") {
|
||||
t.Error("tunnel body must preserve stream:true")
|
||||
}
|
||||
|
||||
// PrepareRun must NOT be called on the tunnel path.
|
||||
if fake.poolPrepareRunCalledSnapshot() {
|
||||
t.Error("PrepareRun must NOT be called on tunnel path")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolStrictOutputNormalizesAgentResponse verifies that the
|
||||
// provider-pool normalized path preserves the strict-output contract from
|
||||
// PrepareRun into completeResponse. A prompt containing an XML completion
|
||||
// contract must have its response wrapped in the expected block (SDD D02).
|
||||
func TestResponsesProviderPoolStrictOutputNormalizesAgentResponse(t *testing.T) {
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathNormalized),
|
||||
}
|
||||
fake.poolRunFrames = make(chan *iop.RunEvent, 3)
|
||||
fake.poolRunFrames <- &iop.RunEvent{Type: "reasoning_delta", Delta: "hidden"}
|
||||
fake.poolRunFrames <- &iop.RunEvent{Type: "delta", Delta: "plain answer"}
|
||||
fake.poolRunFrames <- &iop.RunEvent{Type: "complete", RunId: "run-pool-normalized"}
|
||||
close(fake.poolRunFrames)
|
||||
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{ID: "ollama-strict-model", Providers: map[string]string{"prov-ollama": "served-ollama"}},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{StrictOutput: true}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"ollama-strict-model",
|
||||
"input":"Once you've completed the user's task, you must use the attempt_completion tool to present the result.\n\n<attempt_completion>\n<result>done</result>\n</attempt_completion>\n\nfinish"
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// PrepareRun must have been invoked (normalized path).
|
||||
if !fake.poolPrepareRunCalledSnapshot() {
|
||||
t.Fatal("PrepareRun was not called on normalized path")
|
||||
}
|
||||
|
||||
// The prepared request must carry strict_output and the injected contract
|
||||
// instruction.
|
||||
if len(fake.reqsSnapshot()) != 1 {
|
||||
t.Fatalf("expected one prepared req, got %d", len(fake.reqsSnapshot()))
|
||||
}
|
||||
reqPrep := fake.reqsSnapshot()[0]
|
||||
if reqPrep.Metadata["strict_output"] != "true" {
|
||||
t.Fatalf("strict_output metadata lost in prepared request: %+v", reqPrep.Metadata)
|
||||
}
|
||||
if !strings.Contains(reqPrep.Prompt, "output exactly one <attempt_completion> block") {
|
||||
t.Fatalf("strict contract instruction not injected into prepared prompt:\n%s", reqPrep.Prompt)
|
||||
}
|
||||
|
||||
var resp responsesResponse
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
want := "<attempt_completion>\n<result>plain answer</result>\n</attempt_completion>"
|
||||
if resp.OutputText != want {
|
||||
t.Fatalf("output_text:\ngot %q\nwant %q", resp.OutputText, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResponsesProviderPoolTunnelMetadataAndContextPropagation verifies that the
|
||||
// provider-pool tunnel path preserves metadata, estimated_input_tokens, and
|
||||
// context_class from the base request into the tunnel dispatch (SDD S02/S03).
|
||||
func TestResponsesProviderPoolTunnelMetadataAndContextPropagation(t *testing.T) {
|
||||
frames := make(chan *iop.ProviderTunnelFrame, 3)
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(`{"ok":true}`)}
|
||||
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
||||
close(frames)
|
||||
|
||||
fake := &fakeRunService{
|
||||
poolDispatchPath: string(edgeservice.ProviderPoolPathTunnel),
|
||||
tunnelFrames: frames,
|
||||
tunnelServedTarget: "served-model",
|
||||
}
|
||||
|
||||
catalog := []config.ModelCatalogEntry{
|
||||
{ID: "pool-model", Providers: map[string]string{"prov-1": "served-model"}},
|
||||
}
|
||||
srv := NewServer(config.EdgeOpenAIConf{}, fake, nil)
|
||||
srv.SetModelCatalog(catalog)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{
|
||||
"model":"pool-model",
|
||||
"input":"hello",
|
||||
"metadata":{"experiment":"pool-test"}
|
||||
}`))
|
||||
w := httptest.NewRecorder()
|
||||
srv.handleResponses(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
tunnelReqs := fake.tunnelReqsSnapshot()
|
||||
if len(tunnelReqs) != 1 {
|
||||
t.Fatalf("expected one tunnel req, got %d", len(tunnelReqs))
|
||||
}
|
||||
tunnelReq := tunnelReqs[0]
|
||||
|
||||
// Metadata must contain the principal and pool-echoed fields.
|
||||
if tunnelReq.Metadata == nil {
|
||||
t.Fatalf("tunnel metadata must not be nil")
|
||||
}
|
||||
if tunnelReq.Metadata["openai_model"] == "" {
|
||||
t.Error("openai_model missing from tunnel metadata")
|
||||
}
|
||||
if tunnelReq.Metadata["openai_stream"] == "" {
|
||||
t.Error("openai_stream missing from tunnel metadata")
|
||||
}
|
||||
if tunnelReq.Metadata["estimated_input_tokens"] == "" {
|
||||
t.Error("estimated_input_tokens missing from tunnel metadata")
|
||||
}
|
||||
if tunnelReq.Metadata["context_class"] == "" {
|
||||
t.Error("context_class missing from tunnel metadata")
|
||||
}
|
||||
if tunnelReq.Metadata["experiment"] != "pool-test" {
|
||||
t.Errorf("caller metadata lost: got %v", tunnelReq.Metadata)
|
||||
}
|
||||
|
||||
// EstimatedInputTokens and ContextClass must be non-zero/non-empty on the
|
||||
// tunnel request itself, so Node observability preserves the same values
|
||||
// as the direct tunnel path.
|
||||
if tunnelReq.EstimatedInputTokens <= 0 {
|
||||
t.Errorf("EstimatedInputTokens must be > 0: got %d", tunnelReq.EstimatedInputTokens)
|
||||
}
|
||||
if tunnelReq.ContextClass == "" {
|
||||
t.Error("ContextClass must be non-empty on tunnel request")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -921,6 +921,15 @@ const (
|
|||
// after wire dispatch for provider-pool tunnel paths.
|
||||
type prepareTunnelFunc func(req SubmitProviderTunnelRequest) (SubmitProviderTunnelRequest, error)
|
||||
|
||||
// PrepareRun is an optional pre-dispatch hook that lets the caller prepare
|
||||
// or validate the Run request for the normalized execution path. It runs
|
||||
// AFTER a candidate is selected and ONLY on the normalized path (not on
|
||||
// tunnel/passthrough). If PrepareRun returns an error, the slot is released
|
||||
// and no normalized run is sent. This lets the caller enforce that a complete
|
||||
// SubmitRunRequest (including ModelGroupKey and other required fields) is
|
||||
// present before dispatch, while leaving the tunnel branch untouched.
|
||||
type prepareRunFunc func(req SubmitRunRequest) (SubmitRunRequest, error)
|
||||
|
||||
// ProviderPoolDispatchRequest bundles the Run and Tunnel surface values for
|
||||
// a single one-shot provider-pool dispatch. SubmitProviderPool uses exactly
|
||||
// one queue admission to select a candidate, then dispatches only the
|
||||
|
|
@ -929,6 +938,7 @@ type ProviderPoolDispatchRequest struct {
|
|||
Run SubmitRunRequest
|
||||
Tunnel SubmitProviderTunnelRequest
|
||||
PrepareTunnel prepareTunnelFunc
|
||||
PrepareRun prepareRunFunc
|
||||
}
|
||||
|
||||
// ProviderPoolDispatchResult describes which execution path was selected and
|
||||
|
|
@ -987,6 +997,7 @@ func (s *Service) SubmitProviderPool(ctx context.Context, req ProviderPoolDispat
|
|||
tunnelReq.SessionID = req.Run.SessionID
|
||||
tunnelReq.MaxQueue = req.Run.MaxQueue
|
||||
tunnelReq.QueueTimeoutMS = req.Run.QueueTimeoutMS
|
||||
tunnelReq.Metadata = req.Run.Metadata
|
||||
tunnelReq.EstimatedInputTokens = req.Run.EstimatedInputTokens
|
||||
tunnelReq.ContextClass = req.Run.ContextClass
|
||||
|
||||
|
|
@ -1036,7 +1047,15 @@ func (s *Service) SubmitProviderPool(ctx context.Context, req ProviderPoolDispat
|
|||
}, nil
|
||||
|
||||
case providerExecutionPathNormalized:
|
||||
return s.dispatchProviderPoolRun(ctx, req.Run, adapter, target, selected, queueReason, longReserved)
|
||||
runReq := req.Run
|
||||
if req.PrepareRun != nil {
|
||||
runReq, err = req.PrepareRun(runReq)
|
||||
if err != nil {
|
||||
s.queue.releaseSlotWithLong(req.Run.ModelGroupKey, selected.entry.NodeID, selected.providerID, longReserved)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return s.dispatchProviderPoolRun(ctx, runReq, adapter, target, selected, queueReason, longReserved)
|
||||
|
||||
default:
|
||||
s.queue.releaseSlotWithLong(req.Run.ModelGroupKey, selected.entry.NodeID, selected.providerID, longReserved)
|
||||
|
|
|
|||
Loading…
Reference in a new issue