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

package service

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"iop/packages/go/config"
)

// 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)

// prepareProtocolTunnelFunc prepares a tunnel from the immutable candidate
// selected by the provider pool. It is preferred over PrepareTunnel when both
// hooks are supplied; the legacy hook remains source-compatible.
type prepareProtocolTunnelFunc func(req SubmitProviderTunnelRequest, selected ProviderPoolCandidate) (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
	ProfileID             string
	ProfileDriver         string
	ProfileCapabilities   []string
	ProtocolProfile       *config.ConcreteProtocolProfile
}

// 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")

// ProviderPoolOperationUnsupportedError reports that provider pool candidates
// were rejected because no candidate's concrete protocol profile supports the
// requested operation.
type ProviderPoolOperationUnsupportedError struct {
	Operation string
}

func (e *ProviderPoolOperationUnsupportedError) Error() string {
	return fmt.Sprintf("provider pool does not support operation %q", e.Operation)
}

func (e *ProviderPoolOperationUnsupportedError) Unwrap() error {
	return ErrProviderPoolCandidateRejected
}

// 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.
//
// AvoidProviderID is a request-local recovery hint. When non-empty, every
// admission (initial and queued re-resolution) prefers a runtime-eligible
// alternate provider over the avoided one. The avoided provider is only
// retained when no alternate exists AND AllowAvoidedProviderFallback is
// true AND the provider is still runtime eligible — the explicit fallback
// permission is the only way to re-select the avoided provider. Liveness
// recovery never grants that permission: a health probe proves endpoint
// availability, not safety of replaying the request that just stalled.
//
// Zero values (empty AvoidProviderID, false AllowAvoidedProviderFallback)
// preserve the current candidate selection behavior.
//
// This is selection policy only: it does not create a retry loop, reserve
// a slot, change provider priority, persist the hints, or count retries.
type ProviderPoolDispatchRequest struct {
	Run                          SubmitRunRequest
	Tunnel                       SubmitProviderTunnelRequest
	PrepareProtocolTunnel        prepareProtocolTunnelFunc
	PrepareTunnel                prepareTunnelFunc
	PrepareRun                   prepareRunFunc
	AcceptCandidate              ProviderPoolCandidatePredicate
	AvoidProviderID              string
	AllowAvoidedProviderFallback bool
}

// 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")
	}

	operation := strings.TrimSpace(req.Tunnel.Operation)
	operationPredicate := providerOperationCandidatePredicate(operation)
	candidates, returnedPolicy, err := s.resolveQueueCandidates(req.Run)
	if err != nil {
		return nil, err
	}
	if operationPredicate != nil {
		var rejected bool
		candidates, rejected = filterProviderPoolCandidates(candidates, operationPredicate)
		if rejected {
			return nil, &ProviderPoolOperationUnsupportedError{Operation: operation}
		}
	}
	if req.AcceptCandidate != nil {
		var rejected bool
		candidates, rejected = filterProviderPoolCandidates(candidates, req.AcceptCandidate)
		if rejected {
			return nil, ErrProviderPoolCandidateRejected
		}
	}

	// Request-local avoided-provider recovery preference. This is selection policy
	// only: it does not create a retry loop, reserve a slot, change provider
	// priority, persist the hints, or count retries. The queue owns application —
	// it applies the preference after runtime-health filtering under its lock for
	// both this immediate admission and every queued re-resolution — so a
	// zero-value policy preserves the current candidate set.
	recovery := recoveryCandidatePolicy{
		avoidProviderID:              req.AvoidProviderID,
		allowAvoidedProviderFallback: req.AllowAvoidedProviderFallback,
	}

	// 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 operationPredicate != nil || req.AcceptCandidate != nil {
		resolveCandidates = func() ([]candidateNode, error) {
			resolved, _, err := s.resolveQueueCandidates(req.Run)
			if err != nil {
				return nil, err
			}
			if operationPredicate != nil {
				var rejected bool
				resolved, rejected = filterProviderPoolCandidates(resolved, operationPredicate)
				if rejected {
					return nil, &ProviderPoolOperationUnsupportedError{Operation: operation}
				}
			}
			if req.AcceptCandidate != nil {
				var rejected bool
				resolved, rejected = filterProviderPoolCandidates(resolved, req.AcceptCandidate)
				if rejected {
					return nil, ErrProviderPoolCandidateRejected
				}
			}
			// The avoided-provider recovery preference is NOT applied here: the
			// queue reapplies it under its lock in resolveQueuedCandidatesLocked,
			// after runtime-health and orphan filtering, so re-resolution honors
			// the same request-local hint against genuinely eligible candidates.
			return resolved, nil
		}
	}
	selected, queueReason, err := s.queue.admitWithRecovery(ctx, req.Run.ModelGroupKey, req.Run.Adapter, req.Run.Target, candidates, policy, resolveCandidates, long, req.Run.ProviderPool, recovery)
	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)

	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(ctx, 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
			}
		}
		runReq.ResponseStallTimeoutMS = selected.responseStallTimeoutMS
		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)
	}
}

// providerOperationCandidatePredicate admits legacy candidates without a
// profile, while requiring a concrete profile to declare the requested
// operation. The same predicate is used for initial and queued resolution.
func providerOperationCandidatePredicate(operation string) ProviderPoolCandidatePredicate {
	operation = strings.TrimSpace(operation)
	if operation == "" {
		return nil
	}
	return func(candidate ProviderPoolCandidate) bool {
		if candidate.ProtocolProfile == nil {
			return true
		}
		_, supported := candidate.ProtocolProfile.Operations[operation]
		return supported
	}
}

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 {
		if !accept(providerPoolCandidateSnapshot(&candidate)) {
			continue
		}
		accepted = append(accepted, candidate)
	}
	return accepted, len(accepted) == 0
}

func providerPoolCandidateSnapshot(candidate *candidateNode) ProviderPoolCandidate {
	if candidate == nil {
		return ProviderPoolCandidate{}
	}
	snapshot := ProviderPoolCandidate{
		ActualModel:           candidate.servedTarget,
		ProviderID:            candidate.providerID,
		ExecutionPath:         string(candidate.executionPath),
		LifecycleCapabilities: append([]string(nil), candidate.lifecycleCapabilities...),
	}
	if candidate.profile != nil {
		profile := candidate.profile.Clone()
		snapshot.ProfileID = profile.ID
		snapshot.ProfileDriver = string(profile.Driver)
		snapshot.ProfileCapabilities = append([]string(nil), profile.Capabilities...)
		snapshot.ProtocolProfile = &profile
	}
	return snapshot
}

// profileFacts extracts the immutable profile id and driver from a candidate's
// concrete profile snapshot. Returns empty strings when the profile is nil
// (legacy candidates without a profile).
func profileFacts(p *config.ConcreteProtocolProfile) (id, driver string) {
	if p == nil {
		return "", ""
	}
	return p.ID, string(p.Driver)
}

// 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(
	ctx context.Context,
	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 = selected.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
	if req.PrepareProtocolTunnel != nil {
		tunnelReqPrepared, prepErr := req.PrepareProtocolTunnel(tunnelReq, providerPoolCandidateSnapshot(selected))
		if prepErr != nil {
			reservation.release("prepare-protocol-tunnel-error")
			return nil, prepErr
		}
		tunnelReq = tunnelReqPrepared
	} else if req.PrepareTunnel != nil {
		tunnelReqPrepared, prepErr := req.PrepareTunnel(tunnelReq)
		if prepErr != nil {
			reservation.release("prepare-tunnel-error")
			return nil, prepErr
		}
		tunnelReq = tunnelReqPrepared
	}
	tunnelReq.ResponseStallTimeoutMS = selected.responseStallTimeoutMS

	tunnelReqResolved, runID, err := buildProviderTunnelRequest(tunnelReq, adapter, target)
	if err != nil {
		reservation.release("build-error")
		return nil, err
	}
	if err := s.attachCredentialLease(ctx, tunnelReq, selected.entry, target, tunnelReqResolved); err != nil {
		reservation.release("credential-lease-error")
		return nil, err
	}
	if tunnelReq.CredentialBinding != nil {
		defer s.releaseCredentialLease()
	}

	reservation.track(runID)

	var handle *ProviderTunnelHandle
	err = s.registry.WithCurrentDispatchOwner(selected.entry.NodeID, selected.entry.Client, selected.generation, func() error {
		if err := s.validateCredentialFence(tunnelReqResolved.GetCredentialBinding()); err != nil {
			return err
		}
		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
	disp.ProfileID, disp.ProfileDriver = profileFacts(selected.profile)
	disp.ProfileOperation = tunnelReq.Operation
	if selected.profile != nil {
		disp.ProfileCapabilities = append([]string(nil), selected.profile.Capabilities...)
	}
	handle.RunDispatch = disp
	handle.RunDispatch.ProfileCapabilities = append([]string(nil), disp.ProfileCapabilities...)
	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
	req.ResponseStallTimeoutMS = selected.responseStallTimeoutMS

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

	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()),
		ResponseStallTimeoutMS: dispatchResponseStallTimeout(runReq.GetResponseStallTimeoutMs()),
		EstimatedInputTokens:   req.EstimatedInputTokens,
		ContextClass:           req.ContextClass,
		ProviderID:             selected.providerID,
		UsageAttribution:       req.UsageAttribution,
		ProviderType:           selected.providerType,
		ExecutionPath:          string(selected.executionPath),
		QueueReason:            queueReason,
	}
	disp.ProfileID, disp.ProfileDriver = profileFacts(selected.profile)
	if selected.profile != nil {
		disp.ProfileCapabilities = append([]string(nil), selected.profile.Capabilities...)
	}

	handleDispatch := disp
	handleDispatch.ProfileCapabilities = append([]string(nil), disp.ProfileCapabilities...)
	return &ProviderPoolDispatchResult{
		Path:         ProviderPoolPathNormalized,
		Run:          newRunHandle(handleDispatch, sub),
		DispatchInfo: disp,
	}, nil
}