iop/apps/edge/internal/service/provider_pool.go

package service

import (
	"context"
	"errors"
	"fmt"
)

// providerPoolPath indicates which execution path was selected for a
// provider-pool dispatch. These constants mirror the candidateNode.executionPath
// values but are exposed at the dispatch surface so callers (e.g. OpenAI
// handler) can distinguish tunnel/passthrough from normalized execution.
type providerPoolPath string

const (
	// ProviderPoolPathTunnel means the selected candidate executes via
	// provider tunnel / OpenAI-compatible passthrough.
	ProviderPoolPathTunnel providerPoolPath = "provider_tunnel"
	// ProviderPoolPathNormalized means the selected candidate executes via
	// normalized RunEvent path (Ollama/CLI/native).
	ProviderPoolPathNormalized providerPoolPath = "normalized"
)

// PrepareTunnel is an optional pre-dispatch hook that lets the caller
// inject or modify headers on a tunnel-path request before buildProviderTunnelRequest
// and the Node Send step run. If PrepareTunnel returns an error, the slot is
// released and no tunnel request is sent. This avoids late header mutation
// 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)

// ProviderPoolCandidate is the stable, caller-neutral view passed to a
// request-local provider-pool admission predicate. It intentionally contains
// only actual target and configured provider capability facts; HTTP caller or
// agent identity is never part of pool selection.
type ProviderPoolCandidate struct {
	ActualModel           string
	ProviderID            string
	ExecutionPath         string
	LifecycleCapabilities []string
}

// ProviderPoolCandidatePredicate decides whether a resolved provider candidate
// can serve one request. The service invokes it for both the first admission
// and every queue/recovery re-resolution.
type ProviderPoolCandidatePredicate func(ProviderPoolCandidate) bool

// ErrProviderPoolCandidateRejected reports that otherwise available provider
// candidates were all rejected by a request-local admission predicate before a
// slot was reserved or a provider dispatch was sent.
var ErrProviderPoolCandidateRejected = errors.New("provider pool candidates rejected by request admission policy")

// 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
// execution path indicated by the candidate's executionPath.
type ProviderPoolDispatchRequest struct {
	Run             SubmitRunRequest
	Tunnel          SubmitProviderTunnelRequest
	PrepareTunnel   prepareTunnelFunc
	PrepareRun      prepareRunFunc
	AcceptCandidate ProviderPoolCandidatePredicate
}

// ProviderPoolDispatchResult describes which execution path was selected and
// carries the corresponding dispatch result. Exactly one of Run or Tunnel is
// non-nil.
type ProviderPoolDispatchResult struct {
	Path         providerPoolPath
	Run          RunResult
	Tunnel       ProviderTunnelResult
	DispatchInfo RunDispatch
}

// SubmitProviderPool is the one-shot provider-pool dispatch surface. It performs
// a single queue admission, selects one candidate from the catalog, and dispatches
// the selected execution path — either tunnel/passthrough (OpenAI-compatible
// providers) or normalized RunEvent (Ollama/CLI/native). This method replaces
// the caller's need to choose between SubmitRun and SubmitProviderTunnel.
func (s *Service) SubmitProviderPool(ctx context.Context, req ProviderPoolDispatchRequest) (*ProviderPoolDispatchResult, error) {
	if s.queue == nil {
		return nil, fmt.Errorf("model queue is not configured")
	}

	candidates, returnedPolicy, err := s.resolveQueueCandidates(req.Run)
	if err != nil {
		return nil, err
	}
	candidates, rejected := filterProviderPoolCandidates(candidates, req.AcceptCandidate)
	if rejected {
		return nil, ErrProviderPoolCandidateRejected
	}

	// Provider-pool dispatch uses the canonical policy from the runtime snapshot.
	var policy groupPolicy
	if req.Run.ProviderPool {
		_, _, policy = s.runtimeConfigSnapshot()
	} else {
		policy = returnedPolicy
	}

	long := req.Run.ContextClass == contextClassLong
	resolveCandidates := s.resolveQueueCandidatesClosure(req.Run)
	if req.AcceptCandidate != nil {
		resolveCandidates = func() ([]candidateNode, error) {
			resolved, _, err := s.resolveQueueCandidates(req.Run)
			if err != nil {
				return nil, err
			}
			accepted, rejected := filterProviderPoolCandidates(resolved, req.AcceptCandidate)
			if rejected {
				return nil, ErrProviderPoolCandidateRejected
			}
			return accepted, nil
		}
	}
	selected, queueReason, err := s.queue.admitWithReason(ctx, req.Run.ModelGroupKey, req.Run.Adapter, req.Run.Target, candidates, policy, resolveCandidates, long, req.Run.ProviderPool)
	if err != nil {
		return nil, err
	}

	// The admitted slot is owned by one reservation from here on: every failure
	// path below releases through it, and a dispatched request hands it off to
	// the run/tunnel lifecycle.
	reservation := newQueueReservation(s.queue, selected)

	// Rewrite adapter and target for provider-pool dispatch: the winning candidate
	// carries the concrete adapter and served model name determined at selection time.
	adapter := req.Run.Adapter
	if selected.adapter != "" {
		adapter = selected.adapter
	}
	target := req.Run.Target
	if selected.servedTarget != "" {
		target = selected.servedTarget
	}

	switch selected.executionPath {
	case providerExecutionPathTunnel:
		return s.dispatchProviderPoolTunnel(req, adapter, target, selected, queueReason, reservation)

	case providerExecutionPathNormalized:
		runReq := req.Run
		if req.PrepareRun != nil {
			runReq, err = req.PrepareRun(runReq)
			if err != nil {
				reservation.release("prepare-run-error")
				return nil, err
			}
		}
		return s.dispatchProviderPoolRun(ctx, runReq, adapter, target, selected, queueReason, reservation)

	default:
		reservation.release("unknown-execution-path")
		return nil, fmt.Errorf("unknown execution path %q for provider-pool dispatch", selected.executionPath)
	}
}

func filterProviderPoolCandidates(candidates []candidateNode, accept ProviderPoolCandidatePredicate) ([]candidateNode, bool) {
	if accept == nil || len(candidates) == 0 {
		return candidates, false
	}
	accepted := make([]candidateNode, 0, len(candidates))
	for _, candidate := range candidates {
		capabilities := append([]string(nil), candidate.lifecycleCapabilities...)
		if !accept(ProviderPoolCandidate{
			ActualModel:           candidate.servedTarget,
			ProviderID:            candidate.providerID,
			ExecutionPath:         string(candidate.executionPath),
			LifecycleCapabilities: capabilities,
		}) {
			continue
		}
		accepted = append(accepted, candidate)
	}
	return accepted, len(accepted) == 0
}

// dispatchProviderPoolTunnel relays the selected candidate's raw provider
// request after provider-pool admission. The tunnel inherits the Run's
// identity, metadata, and long-context classification so passthrough dispatch
// keeps the same observability as the normalized path.
func (s *Service) dispatchProviderPoolTunnel(
	req ProviderPoolDispatchRequest,
	adapter, target string,
	selected *candidateNode,
	queueReason string,
	reservation *queueReservation,
) (*ProviderPoolDispatchResult, error) {
	tunnelReq := req.Tunnel
	tunnelReq.ProviderPool = true
	tunnelReq.ModelGroupKey = req.Run.ModelGroupKey
	tunnelReq.ProviderID = req.Run.ProviderID
	tunnelReq.UsageAttribution = req.Run.UsageAttribution
	tunnelReq.Adapter = adapter
	tunnelReq.Target = target
	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

	// Apply pre-dispatch tunnel preparation (e.g. provider auth headers)
	// before buildProviderTunnelRequest so headers reach the wire request.
	if req.PrepareTunnel != nil {
		tunnelReqPrepared, prepErr := req.PrepareTunnel(tunnelReq)
		if prepErr != nil {
			reservation.release("prepare-tunnel-error")
			return nil, prepErr
		}
		tunnelReq = tunnelReqPrepared
	}

	tunnelReqResolved, runID, err := buildProviderTunnelRequest(tunnelReq, adapter, target)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}

	reservation.track(runID)

	var handle *ProviderTunnelHandle
	err = s.registry.WithCurrentDispatchOwner(selected.entry.NodeID, selected.entry.Client, selected.generation, func() error {
		h, err := s.openProviderTunnel(selected.entry, tunnelReqResolved, tunnelReq, queueReason, true, selected.providerID, selected.providerType, string(selected.executionPath))
		if err != nil {
			return err
		}
		handle = h
		reservation.handOff()
		return nil
	})
	if err != nil {
		if !s.candidateIsCurrentOwner(selected) {
			reservation.release("stale-generation")
			return nil, staleGenerationError(selected)
		}
		reservation.release("send-error")
		return nil, err
	}

	disp := handle.Dispatch()
	disp.ProviderID = selected.providerID
	disp.UsageAttribution = req.Run.UsageAttribution
	disp.ProviderType = selected.providerType
	disp.ExecutionPath = string(selected.executionPath)
	disp.QueueReason = queueReason
	return &ProviderPoolDispatchResult{
		Path:         ProviderPoolPathTunnel,
		Tunnel:       handle,
		DispatchInfo: disp,
	}, nil
}

// dispatchProviderPoolRun is a helper that submits a normalized RunRequest after
// provider-pool admission. It tracks inflight, sends the RunRequest to the
// selected entry, and returns a RunResult wrapping the dispatch info.
func (s *Service) dispatchProviderPoolRun(
	ctx context.Context,
	req SubmitRunRequest,
	adapter, target string,
	selected *candidateNode,
	queueReason string,
	reservation *queueReservation,
) (*ProviderPoolDispatchResult, error) {
	req.Adapter = adapter
	req.Target = target

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}

	// Track inflight before send so the event watcher can release the slot even
	// if a terminal event arrives before the Send call completes.
	reservation.track(runID)

	var sub *runSubscription
	var subErr error
	err = s.registry.WithCurrentDispatchOwner(selected.entry.NodeID, selected.entry.Client, selected.generation, func() error {
		sb, err := s.subscribeRun(runID, selected.entry.NodeID, runReq.GetBackground())
		if err != nil {
			subErr = err
			return err
		}
		if err := selected.entry.Client.Send(runReq); err != nil {
			sb.close()
			return err
		}
		sub = sb
		reservation.handOff()
		return nil
	})
	if err != nil {
		if !s.candidateIsCurrentOwner(selected) {
			reservation.release("stale-generation")
			return nil, staleGenerationError(selected)
		}
		if subErr != nil {
			reservation.release("no-event-bus")
			return nil, subErr
		}
		reservation.release("send-error")
		return nil, err
	}

	disp := RunDispatch{
		RunID:                runID,
		NodeID:               selected.entry.NodeID,
		NodeLabel:            nodeLabel(selected.entry),
		ModelGroupKey:        req.ModelGroupKey,
		Adapter:              runReq.GetAdapter(),
		Target:               runReq.GetTarget(),
		SessionID:            runReq.GetSessionId(),
		Background:           runReq.GetBackground(),
		TimeoutSec:           int(runReq.GetTimeoutSec()),
		EstimatedInputTokens: req.EstimatedInputTokens,
		ContextClass:         req.ContextClass,
		ProviderID:           selected.providerID,
		UsageAttribution:     req.UsageAttribution,
		ProviderType:         selected.providerType,
		ExecutionPath:        string(selected.executionPath),
		QueueReason:          queueReason,
	}

	return &ProviderPoolDispatchResult{
		Path:         ProviderPoolPathNormalized,
		Run:          newRunHandle(disp, sub),
		DispatchInfo: disp,
	}, nil
}